Update inference.py

inference.py

@@ -21,11 +21,12 @@ def load_model_and_tokenizer():
     model = AutoModelForCausalLM.from_pretrained(
     MODEL_REPO,
     trust_remote_code=True,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
+    torch_dtype=torch.bfloat16 if torch.cuda.is_available() else torch.float32,
     device_map="auto"
 )
 
 
+
     return tokenizer, model