0.14 explictly set GPU and dtype

app.py

@@ -24,11 +24,11 @@ models_available = [
     "mistralai/Mistral-7B-Instruct-v0.3",
 ]
 
-device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+#device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+device = "cuda" 
 
 tokenizer_a, model_a = None, None
 tokenizer_b, model_b = None, None
-torch_dtype = torch.bfloat16
 
 def apply_chat_template(messages, add_generation_prompt=False):
     """
@@ -62,7 +62,7 @@ def load_model_a(model_id):
         tokenizer_a = AutoTokenizer.from_pretrained(model_id)
         model_a = AutoModelForCausalLM.from_pretrained(
             model_id,
-            torch_dtype=torch_dtype,
+            torch_dtype=torch.bfloat16,
             device_map="auto",
             trust_remote_code=True,
         ).eval()
@@ -79,7 +79,7 @@ def load_model_b(model_id):
         tokenizer_b = AutoTokenizer.from_pretrained(model_id)
         model_b = AutoModelForCausalLM.from_pretrained(
             model_id,
-            torch_dtype=torch_dtype,
+            torch_dtype=torch.bfloat16,
             device_map="auto",
             trust_remote_code=True,
         ).eval()