Spaces:

timm
/

timmAttentionViz

Running

rwightman HF staff commited on Aug 29

Commit

e9cfdba

•

1 Parent(s): 6458094

Update app.py

Files changed (1) hide show

app.py CHANGED Viewed

@@ -84,8 +84,6 @@ def rollout(attentions, discard_ratio, head_fusion, num_prefix_tokens=1):
             # Discard the lowest attentions, but don't discard the prefix tokens
             flat = attention_heads_fused.view(-1)
             _, indices = flat.topk(int(flat.size(-1 )* discard_ratio), -1, False)
-            print(indices)
-            print(indices.shape)
             indices = indices[indices >= num_prefix_tokens]
             flat[indices] = 0

             # Discard the lowest attentions, but don't discard the prefix tokens
             flat = attention_heads_fused.view(-1)
             _, indices = flat.topk(int(flat.size(-1 )* discard_ratio), -1, False)
             indices = indices[indices >= num_prefix_tokens]
             flat[indices] = 0