Update app.py

app.py

@@ -99,7 +99,7 @@ def get_hidden_states(raw_original_prompt, force_hidden_states=False):
         important_tokens = (1 + 
                             F.normalize(hidden_states, dim=-1).diff(dim=0)
                             .norm(dim=-1).topk(k=4).indices[:, 0].cpu().numpy())
-        print(f'{important_tokens=}')
+        print(f'{important_tokens=}\t\t{hidden_states.shape=}')
         global_state.local_state.hidden_states = hidden_states.cpu().detach()
         
     token_btns = ([gr.Button(token, visible=True,