Update dependencyvit.py

Author: fffffgggg54

timm/models/dependencyvit.py

@@ -124,8 +124,8 @@ def forward(self, in_tuple: Tuple[torch.Tensor, torch.Tensor]) -> Tuple[torch.Te
         #prune_mask = attn.detach().abs().sum((1, -1))
         #prune_mask = attn.sum(1).sum(-1)
         #prune_mask = attn.sum(1).abs().sum(-1)
-        prune_mask = attn.abs().sum((1, -1))
-        #prune_mask = m.reshape(B, N)
+        #prune_mask = attn.abs().sum((1, -1))
+        prune_mask = m.reshape(B, N)
 
         x = self.proj(x)
         x = self.proj_drop(x)