Update darnabot.py

darnabot.py

@@ -72,15 +72,16 @@ import spaces
 MODEL_LIST = ["mistralai/Mistral-Nemo-Instruct-2407"]
 HF_TOKEN = os.environ.get("HF_TOKEN", None)
 MODEL = os.environ.get("MODEL_ID")
-device = "cuda" if torch.cuda.is_available() else "cpu"  # Use GPU if available, otherwise CPU
+
+device = "cuda"
 
 tokenizer = AutoTokenizer.from_pretrained(MODEL)
 model = AutoModelForCausalLM.from_pretrained(
     MODEL,
-    torch_dtype=torch.bfloat16,  # Using bf16 precision
-    device_map="auto",  # Automatically map model to available devices
-    ignore_mismatched_sizes=True  # Allows some flexibility if dimensions don't match
-).to(device)
+    torch_dtype=torch.bfloat16,
+    device_map="auto",
+    ignore_mismatched_sizes=True)
+
 
 async def chat(messages):
     # Convert messages to the format required for the model