Update app.py

app.py

@@ -55,7 +55,7 @@ def run_lens(model,PROMPT):
             logits_lens_token_result_by_layer.append(logits_lens_next_token)
         tokens_out = llama.lm_head.output.argmax(dim=-1).save()
         expected_token = tokens_out[0][-1].save()
-    logits_lens_all_probs = np.concatenate([probs[:, expected_token].cpu().detach().numpy() for probs in logits_lens_probs_by_layer])
+    logits_lens_all_probs = np.concatenate([probs[:, expected_token].cpu().detach().to(torch.float32).numpy() for probs in logits_lens_probs_by_layer])
     #get the rank of the expected token from each layer's distribution
     for layer_probs in logits_lens_probs_by_layer:
         # Sort the probabilities in descending order and find the rank of the expected token