Update app.py

app.py

@@ -56,7 +56,8 @@ def generate(
     if input_ids.shape[1] > MAX_INPUT_TOKEN_LENGTH:
         input_ids = input_ids[:, -MAX_INPUT_TOKEN_LENGTH:]
         gr.Warning(f"Trimmed input from conversation as it was longer than {MAX_INPUT_TOKEN_LENGTH} tokens.")
-    input_ids = input_ids.to(model.device)
+    
+    input_ids = input_ids.to(device)  # Ensure the input tensor is on the correct device
 
     outputs = []
     generated_ids = model.generate(
@@ -69,9 +70,13 @@ def generate(
         num_beams=1,
         repetition_penalty=repetition_penalty
     )
+
+    generated_ids = generated_ids.to(device)  # Ensure the generated ids are moved to the device
+
     outputs.append(tokenizer.decode(generated_ids[0], skip_special_tokens=True))
     return "".join(outputs)
 
+
 chat_interface = gr.ChatInterface(
     fn=generate,
     additional_inputs=[