Update app.py

app.py

@@ -81,7 +81,7 @@ def get_hidden_states(raw_original_prompt):
                   + [gr.Button('', visible=False) for _ in range(MAX_PROMPT_TOKENS - len(tokens))])
     progress_dummy_output = ''
     invisible_bubbles = [gr.Textbox('', visible=False) for i in range(MAX_NUM_LAYERS)]
-    global_state.hidden_states = hidden_states
+    global_state.hidden_states = hidden_states.cpu()
     return [progress_dummy_output, *token_btns, *invisible_bubbles]