change script

scripts/generate_scripts.py

@@ -14,9 +14,10 @@ from transformers import AutoModelForCausalLM, AutoTokenizer
 
 
 @spaces.GPU()
-def generate_local(model,tokenizer, prompt: str, max_new_tokens: int = 350, temperature: float = 0.7) -> str:
-    device = torch.device("cuda" if torch.cuda.is_available() else "cpu")  # get the device the model is on
-    inputs = tokenizer(prompt, return_tensors="pt").to(device)
+def generate_local(model, tokenizer, prompt: str, max_new_tokens: int = 350, temperature: float = 0.7) -> str:
+    inputs = tokenizer(prompt, return_tensors="pt")
+    inputs = {k: v.to(model.device) for k, v in inputs.items()}  # ⬅️ Safely match model's device
+
     output_ids = model.generate(
         **inputs,
         max_new_tokens=max_new_tokens,