Update app.py

app.py

@@ -5,7 +5,7 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 # Load the model and tokenizer locally
 model_name = "kz919/QwQ-0.5B-Distilled-SFT"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name).to("cuda")
 
 # Define the function to handle chat responses
 @spaces.GPU
@@ -17,7 +17,7 @@ def respond(message, history: list[tuple[str, str]], system_message, max_tokens,
     prompt += f"User: {message}\nAssistant:"
 
     # Tokenize the input prompt
-    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = tokenizer(prompt, return_tensors="pt").to("cuda")
 
     # Generate a response
     outputs = model.generate(