fix

app.py

@@ -125,7 +125,7 @@ def chat_with_model(messages):
             output_text += token_str
 
         # Early stopping if user reappears
-        if "\nUser:" in output_text:
+        if "\nUser" in output_text:
             output_text = output_text.split("\nUser")[0].rstrip()
             messages[-1]["content"] = output_text
             break
@@ -145,7 +145,7 @@ def chat_with_model(messages):
         messages[-1]["content"] = output_text
     
     # Wait for thread to finish
-    current_model.to("cpu")
+    # current_model.to("cpu")
     torch.cuda.empty_cache()
 
     messages[-1]["content"] = output_text