Update app.py

app.py

@@ -14,8 +14,8 @@ def load_model():
     model = AutoModelForCausalLM.from_pretrained(
         "umar141/Gemma_1B_Baro_v2_vllm",
         torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto"  # Automatically puts model on CUDA if available
-    )
+
+    ).to("cuda" if torch.cuda.is_available() else "cpu")
     return tokenizer, model
 
 tokenizer, model = load_model()