Spaces:

JeffMII
/

CEC-Learning

Running

Jeff Myers II commited on May 6

Commit

424c50e

1 Parent(s): 47f4077

Completed Prototype

Files changed (1) hide show

Gemma.py CHANGED Viewed

@@ -18,7 +18,7 @@ class GemmaLLM:
         self.model = Gemma3ForCausalLM.from_pretrained(
             model_id,
             device_map="cuda" if torch.cuda.is_available() else "cpu",
-            torch_dtype=torch.float16,
             # token=os.environ.get("GEMMA_TOKEN"),
         ).eval()

         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"),
         ).eval()