Fix model to cuda

app.py

@@ -19,7 +19,7 @@ def load_model_on_selection(model_name, progress=gr.Progress(track_tqdm=False)):
     current_model = AutoModelForCausalLM.from_pretrained(
         model_name,
         torch_dtype=torch.float16,
-        device_map="cuda",
+        device_map="cpu",
         use_auth_token=token
     )
 
@@ -31,7 +31,7 @@ def generate_text(prompt):
     global current_model, current_tokenizer
     if current_model is None or current_tokenizer is None:
         return "⚠️ No model loaded yet. Please select a model first."
-    
+    current_model.to('cuda')
     inputs = current_tokenizer(prompt, return_tensors="pt").to(current_model.device)
     outputs = current_model.generate(**inputs, max_new_tokens=256)
     return current_tokenizer.decode(outputs[0], skip_special_tokens=True)