Wi-zz
/

joy-caption-pre-alpha

Model card Files Files and versions Community

Wi-zz commited on Aug 29

Commit

3bce52e

•

1 Parent(s): 5a6d14f

Update app-multi-alpha.py

Files changed (1) hide show

app-multi-alpha.py +1 -1

app-multi-alpha.py CHANGED Viewed

@@ -78,7 +78,7 @@ def stream_chat(input_images: List[Image.Image], batch_size: int, pbar: tqdm, mo
             print("Skipping this batch and continuing...")
             continue
-        with torch.amp.autocast_mode.autocast(rank, enabled=True):
             vision_outputs = clip_model(pixel_values=images, output_hidden_states=True)
             image_features = vision_outputs.hidden_states[-2]
             embedded_images = image_adapter(image_features).to(dtype=torch.bfloat16)

             print("Skipping this batch and continuing...")
             continue
+        with torch.amp.autocast_mode.autocast(device_type='cuda', enabled=True):
             vision_outputs = clip_model(pixel_values=images, output_hidden_states=True)
             image_features = vision_outputs.hidden_states[-2]
             embedded_images = image_adapter(image_features).to(dtype=torch.bfloat16)