Update app.py

app.py

@@ -26,9 +26,11 @@ def chat(message: str, history: list, stm_state: torch.Tensor, temperature: floa
     model.load_stm_state(stm_state)
     
     response = ""
-    for token_id in model.interact(**tokenized_query, max_seq_len=seq_len, temperature=temperature, top_p=top_p):
-        response += model.stringify_token(token_id, show_memory_update=True)
-        yield history + [[message, response]], stm_state
+
+    with torch.amp.autocast(device_type=device.type, dtype=torch.bfloat16):
+        for token_id in model.interact(**tokenized_query, max_seq_len=seq_len, temperature=temperature, top_p=top_p):
+            response += model.stringify_token(token_id, show_memory_update=True)
+            yield history + [[message, response]], stm_state
     
     return history + [[message, response]], model.export_stm_state().cpu()