Update app.py

app.py

@@ -58,13 +58,10 @@ def generate_and_visualize(prompt, num_tokens=10):
         output_logits = model(inputs_embeds=input_embeds.requires_grad_()).logits
         max_logits, max_indices = torch.max(output_logits[0, -1, :], dim=-1)
 
-        max_logits.backward(max_logits)
-        try:
-            relevance = input_embeds.grad.float().sum(-1).cpu()[0]
-            all_relevances.append(relevance)
-        except:
-            all_relevances.append(0)
-
+        max_logits.backward(max_logits) 
+        relevance = input_embeds.grad.float().sum(-1).cpu()[0]
+        all_relevances.append(relevance)
+        
         next_token = max_indices.unsqueeze(0)
         generated_tokens_ids.append(next_token.item())