Update app.py

app.py

@@ -339,6 +339,11 @@ from langchain_huggingface import HuggingFacePipeline
 from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
 import gradio as gr
 
+import torch
+print(torch.cuda.is_available())  # Check if a GPU is available
+a = torch.Tensor([1]).cuda()  # Should succeed without errors
+print(a)
+
 # Define HuggingFace LLM
 def create_llm():
     model_name = "Qwen/Qwen2.5-7B-Instruct-1M"
@@ -348,7 +353,7 @@ def create_llm():
         task="text-generation",
         model=model,
         tokenizer=tokenizer,
-        device_map="auto",
+        device=0,
         max_new_tokens=500,
     )
     return HuggingFacePipeline(pipeline=llm_pipeline)