Convert pixel_values to BFloat16 before moving to CUDA

app.py

@@ -93,21 +93,14 @@ def get_som_response(instruction, image_som):
         add_generation_prompt=True
     )
 
-    # inputs = magma_processor(images=[image_som], texts=prompt, return_tensors="pt")
-    # inputs['pixel_values'] = inputs['pixel_values'].unsqueeze(0)
-    # inputs['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    # inputs = inputs.to("cuda")
-
     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['image_sizes'] = inputs['image_sizes'].unsqueeze(0)
-    # Convert all tensors to BFloat16 before moving to CUDA
-    for key in inputs:
-        if torch.is_tensor(inputs[key]):
-            inputs[key] = inputs[key].to(torch.bfloat16)
-    # Then move to CUDA device
     inputs = inputs.to("cuda")
 
+
     magam_model.generation_config.pad_token_id = magma_processor.tokenizer.pad_token_id
     with torch.inference_mode():
         output_ids = magam_model.generate(