using what worked the first time

main.py

@@ -33,7 +33,7 @@ tokenizer = AutoTokenizer.from_pretrained(MODEL, token=HF_TOKEN, use_fast=True)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
     token=HF_TOKEN,
-    torch_dtype=torch.float64,  # Use bfloat16 for CPUs, float16 for GPUs
+    torch_dtype=torch.float16,
     device_map="auto"
 )