Completed Prototype

Gemma.py

@@ -18,9 +18,10 @@ class GemmaLLM:
         self.model = Gemma3ForCausalLM.from_pretrained(
             model_id,
             device_map="cuda" if torch.cuda.is_available() else "cpu",
-            torch_dtype=torch.bfloat16,
-            # token=os.environ.get("GEMMA_TOKEN"),
+            torch_dtype=torch.float16,
         ).eval()
+        
+        self.model = self.model.bfloat16()
 
     @spaces.GPU
     def generate(self, message) -> str: