Update app.py

app.py

@@ -84,7 +84,7 @@ def generate(text, past_key_values):
         {'role': 'user', 'content': text}
     ]
     response = pipe(messages, 
-                    past_key_values=past_key_values, #past_kv_to_device(past_key_values, pipe.model.device), 
+                    past_key_values=past_kv_to_device(past_key_values, pipe.model.device), 
                     **generate_kwargs)[0]['generated_text']
     return response[-1]['content']
 
@@ -93,6 +93,6 @@ if __name__ == "__main__":
     with torch.no_grad():
         past_key_values = set_past_key_values()
         print(f'{past_key_values=}')
-        demo = gr.Interface(partial(generate, past_key_values=past_kv_to_device(past_key_values, pipe.model.device)), 
+        demo = gr.Interface(partial(generate, past_key_values=past_key_values), 
                             inputs="textbox", outputs="textbox")
         demo.launch()