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()