Update app.py

app.py

@@ -37,7 +37,7 @@ def initialize_mistral():
     tokenizer = AutoTokenizer.from_pretrained(model_path)
     device = 'cuda' if torch.cuda.is_available() else 'cpu'
     dtype = torch.bfloat16
-    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype, device_map=device)
+    model = AutoModelForCausalLM.from_pretrained(model_path, torch_dtype=dtype)  # Removed device_map parameter
     return tokenizer, model
 
 # Mistral text generation function