Fixed SET negativity pruning. See #19.

TimDettmers · Sep 11, 2020 · f99c2f2 · f99c2f2
1 parent 6b97851
commit f99c2f2
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/sparselearning/funcs.py b/sparselearning/funcs.py
@@ -159,7 +159,7 @@ def magnitude_and_negativity_prune(masking, mask, weight, name):
 
     # remove the most negative weights
     x, idx = torch.sort(weight.data.view(-1))
-    mask.data.view(-1)[idx[:math.ceil(num_remove/2.0)]] = 0.0
+    mask.data.view(-1)[idx[math.ceil(num_remove/2.0):]] = 0.0
 
     return mask