Update app.py

app.py

@@ -92,7 +92,10 @@ def detach_past_kv(past_kv):
 @spaces.GPU
 def set_past_key_values():
     model, tokenizer = pipe.model, pipe.tokenizer
-    tokenized = tokenizer(system_only_prompt_format.format(system_message=system_prompt))
+    tokenized = tokenizer(
+        system_only_prompt_format.format(system_message=system_prompt),
+        return_tensors='pt'
+    )
     # tokenized = tokenizer.apply_chat_template(start_messages, return_tensors='pt')    
     # Check that this is indeed a prefix of the entire message
     # test_messages = [*start_messages, {'role': 'user', 'content': 'Hello World!'}]