Update app.py

app.py

@@ -55,8 +55,12 @@ def initialize_model():
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
         token=token,
-        device_map="auto"
+        device_map="cuda"
     )
+    if torch.cuda.is_available():
+        model.device = "cuda"
+    else:
+        print("CUDA is not available")
     
     return model, tokenizer