Update app.py

app.py

@@ -6,13 +6,8 @@ model_name = "kz919/QwQ-0.5B-Distilled-SFT"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
 model = AutoModelForCausalLM.from_pretrained(model_name)
 
-# Ensure the model runs on GPU if available
-import torch
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-model.to(device)
-
-
 # Define the function to handle chat responses
+@spaces.GPU
 def respond(message, history: list[tuple[str, str]], system_message, max_tokens, temperature, top_p):
     # Prepare the prompt by combining history and system messages
     prompt = system_message + "\n"
@@ -21,7 +16,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").to(device)
+    inputs = tokenizer(prompt, return_tensors="pt")
 
     # Generate a response
     outputs = model.generate(