Apply pruning to DINOv2 meet problems: "IndexError: index 768 is out of bounds for dimension 0 with size 768"

[mr-lz]

Your work is excellent, but the pruning method you proposed results in an error when applied to dinov2: IndexError: index 768 is out of bounds for dimension 0 with size 768.
The error message is as follows:

Traceback (most recent call last):
  File "/search/odin/prune/prune2.py", line 86, in <module>
    for group in pruner.step(interactive=True): # Warning: groups must be handled sequentially. Do not keep them as a list.
  File "/search/odin/anaconda3/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 35, in generator_context
    response = gen.send(None)
  File "/search/odin/anaconda3/lib/python3.9/site-packages/torch_pruning/pruner/algorithms/metapruner.py", line 420, in _prune
    imp = self.estimate_importance(group) # raw importance score
  File "/search/odin/anaconda3/lib/python3.9/site-packages/torch_pruning/pruner/algorithms/metapruner.py", line 282, in estimate_importance
    return self.importance(group)
  File "/search/odin/anaconda3/lib/python3.9/site-packages/torch/utils/_contextlib.py", line 115, in decorate_context
    return func(*args, **kwargs)
  File "/search/odin/anaconda3/lib/python3.9/site-packages/torch_pruning/pruner/importance.py", line 229, in __call__
    local_imp = local_imp[idxs]
IndexError: index 768 is out of bounds for dimension 0 with size 768

Do you have any suggestions?