Update llm.py to set model torch_dtype to 'auto'

2026-05-25 01:58:21 -05:00 · 2025-04-08 09:19:14 -04:00
parent d3e203930c
commit ce27126502
1 changed files with 1 additions and 1 deletions
@@ -16,7 +16,7 @@ class LLM:
        self.tokenizer = AutoTokenizer.from_pretrained(model_name)
        self.model = AutoModelForCausalLM.from_pretrained(
            model_name,
-            torch_dtype=torch.bfloat16,
+            torch_dtype='auto',
            device_map=device,
        ).to(device)