pytorch · fmassa · May 3, 2018 · Apr 24, 2018 · Apr 24, 2018 · Apr 24, 2018
diff --git a/test/test_nn.py b/test/test_nn.py
@@ -5121,7 +5121,7 @@ def test_sparse_default_std(self):
 
                 for col_idx in range(input_tensor.size(1)):
                     column = input_tensor[:, col_idx]
-                    assert column[column == 0].nelement() >= math.ceil(sparsity * cols)
+                    assert column[column == 0].nelement() >= math.ceil(sparsity * rows)
 
                 assert self._is_normal(input_tensor[input_tensor != 0], 0, std)
 

diff --git a/torch/nn/init.py b/torch/nn/init.py
@@ -360,17 +360,14 @@ def sparse_(tensor, sparsity, std=0.01):
         raise ValueError("Only tensors with 2 dimensions are supported")
 
     rows, cols = tensor.shape
-    num_zeros = int(math.ceil(rows * sparsity))
+    num_zeros = int(math.ceil(sparsity * rows))
 
     with torch.no_grad():
         tensor.normal_(0, std)
         for col_idx in range(cols):
-            row_indices = list(range(rows))
-            random.shuffle(row_indices)
+            row_indices = torch.randperm(rows)
             zero_indices = row_indices[:num_zeros]
-            for row_idx in zero_indices:
-                tensor[row_idx, col_idx] = 0
-
+            tensor[zero_indices, col_idx] = 0
     return tensor