Update modeling_internlm2.py

modeling_internlm2.py

@@ -462,9 +462,9 @@ class InternLM2Attention(nn.Module):
             '''
             # compute headmask based on attn weights ratio.
             if attn_weights.shape[2]>1:
-                res_v = torch.sum( attn_weights[0,:,:,41:image_token_num],dim=[1,2])
-                res_t = torch.sum( attn_weights[0,:,:,image_token_num:],dim=[1,2])
-                res_s = torch.sum( attn_weights[0,:,:,:41],dim=[1,2])
+                res_v = torch.sum( attn_weights[0,:,:,self.offset:self.offset+image_token_num],dim=[1,2])
+                res_t = torch.sum( attn_weights[0,:,:,self.offset+image_token_num:],dim=[1,2])
+                res_s = torch.sum( attn_weights[0,:,:,:self.offset],dim=[1,2])
                 res = res_v/(res_t+res_s)
                 torch.save(res, 'headcut_mask_8B/'+str(idx)+'.pth')
                 if idx ==31: