Merge branch 'master' of https://github.com/VainF/Torch-Pruning

VainF · Dec 29, 2023 · e80eaf5 · e80eaf5
2 parents 97a9b11 + aeb9bef
commit e80eaf5
Showing 1 changed file with 3 additions and 5 deletions.
diff --git a/torch_pruning/pruner/importance.py b/torch_pruning/pruner/importance.py
@@ -77,13 +77,11 @@ def __init__(self,
         self.bias = bias
 
     def _lamp(self, imp): # Layer-adaptive Sparsity for the Magnitude-based Pruning
-        argsort_idx = torch.argsort(imp, dim=0, descending=True).tolist()
-        sorted_imp = imp[argsort_idx]
+        argsort_idx = torch.argsort(imp, dim=0, descending=True)
+        sorted_imp = imp[argsort_idx.tolist()]
         cumsum_imp = torch.cumsum(sorted_imp, dim=0)
         sorted_imp = sorted_imp / cumsum_imp
-        inversed_idx = torch.arange(len(sorted_imp))[
-            argsort_idx
-        ].tolist()  # [0, 1, 2, 3, ..., ]
+        inversed_idx = torch.argsort(argsort_idx).tolist()  # [0, 1, 2, 3, ..., ]
         return sorted_imp[inversed_idx]
 
     def _normalize(self, group_importance, normalizer):