Update app.py

app.py

@@ -48,11 +48,11 @@ def initialize_model():
 
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_ID,
-        device_map="auto",
+        device_map="cuda",
         quantization_config=quantization_config,
         torch_dtype=torch.bfloat16,
         trust_remote_code=True
-    )
+    ).to("cuda")
 
     return model, tokenizer