type fixes

app.py

@@ -4,9 +4,13 @@ import torch
 from transformers import AutoModelForImageTextToText, AutoProcessor
 
 # Load model and processor
-MODEL_PATH = "google/gemma-3n-E2B-it"
+MODEL_PATH = "google/gemma-3n-E4B-it"
 processor = AutoProcessor.from_pretrained(MODEL_PATH)
-model = AutoModelForImageTextToText.from_pretrained(MODEL_PATH, torch_dtype="auto", device_map="auto")
+model = AutoModelForImageTextToText.from_pretrained(
+    MODEL_PATH,
+    torch_dtype=torch.bfloat16,
+    device_map="cuda"
+).eval()
 
 @spaces.GPU
 def process_inputs(image, audio):
@@ -18,10 +22,11 @@ def process_inputs(image, audio):
     ).to(model.device, dtype=model.dtype)
 
     # Generate text output
-    outputs = model.generate(
-        **inputs,
-        max_new_tokens=256
-    )
+    with torch.inference_mode:
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=256
+        )
 
     # Decode and return text
     text = processor.batch_decode(