@torch .autocast(device_type=cuda, dtype=torch.bfloat16)

app.py

@@ -95,8 +95,7 @@ def get_som_response(instruction, image_som):
 
     inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
     inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
-    # Convert pixel_values to BFloat16 before moving to CUDA
-    inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
+    # inputs['pixel_values'] = inputs['pixel_values'].to(torch.bfloat16)
     inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
     inputs = inputs.to("cuda")
 
@@ -156,7 +155,7 @@ def get_qa_response(instruction, image):
 
 @spaces.GPU
 @torch.inference_mode()
-# @torch.autocast(device_type="cuda", dtype=torch.bfloat16)
+@torch.autocast(device_type="cuda", dtype=torch.bfloat16)
 def process(
     image_input,
     box_threshold,