jaandoui
/

DNABERT2-AttentionExtracted

Inference Endpoints

Model card Files Files and versions Community

jaandoui commited on May 19

Commit

cb70b22

•

1 Parent(s): b2f170e

Update bert_layers.py

Files changed (1) hide show

bert_layers.py +2 -2

bert_layers.py CHANGED Viewed

@@ -199,9 +199,9 @@ class BertUnpadSelfAttention(nn.Module):
         # print(f'PROBLEM HERE: UNDERSTAND IT!!')
         rearranged_attention = rearrange(attention, 'nnz h d -> nnz (h d)')
-        try:
             # print(f'REARRANGED ATTENTION: {rearranged_attention.shape}')
-        except:
             # print(f'REARRANGED ATTENTION: {rearranged_attention[0].shape}')
         return rearrange(attention, 'nnz h d -> nnz (h d)'), attention_probs

         # print(f'PROBLEM HERE: UNDERSTAND IT!!')
         rearranged_attention = rearrange(attention, 'nnz h d -> nnz (h d)')
+        # try:
             # print(f'REARRANGED ATTENTION: {rearranged_attention.shape}')
+        # except:
             # print(f'REARRANGED ATTENTION: {rearranged_attention[0].shape}')
         return rearrange(attention, 'nnz h d -> nnz (h d)'), attention_probs