use autocast

app.py

@@ -63,18 +63,18 @@ def generate(
         thread = Thread(target=current_model.generate, kwargs=generate_kwargs)
         thread.start()
 
-    # Write the prompt in blue
-    output = "<span style='color: blue;'>" + prompt + "</span>"
-    for new_text in streamer:
-        if isinstance(new_text, torch.Tensor):
-            new_text = current_tokenizer.decode(new_text)
-        if streamer.stop_signal in new_text:
-            output += new_text.split(streamer.stop_signal)[0]
-            break
-        output += new_text
-        yield output
-
-    thread.join()
+        # Write the prompt in blue
+        output = "<span style='color: blue;'>" + prompt + "</span>"
+        for new_text in streamer:
+            if isinstance(new_text, torch.Tensor):
+                new_text = current_tokenizer.decode(new_text)
+            if streamer.stop_signal in new_text:
+                output += new_text.split(streamer.stop_signal)[0]
+                break
+            output += new_text
+            yield output
+
+        thread.join()
     return output
 
 additional_inputs=[