Update app.py

app.py

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