nonzero warnings

Author: rusty1s

torch_sparse/masked_select.py

@@ -65,7 +65,7 @@ def masked_select(src: SparseTensor, dim: int,
     else:
         value = src.storage.value()
         if value is not None:
-            idx = mask.nonzero().flatten()
+            idx = mask.nonzero(as_tuple=False).flatten()
             return src.set_value(value.index_select(dim - 1, idx),
                                  layout='coo')
         else:

torch_sparse/storage.py

@@ -379,7 +379,7 @@ def coalesce(self, reduce: str = "add"):
 
         value = self._value
         if value is not None:
-            ptr = mask.nonzero().flatten()
+            ptr = mask.nonzero(as_tuple=False).flatten()
             ptr = torch.cat([ptr, ptr.new_full((1, ), value.size(0))])
             value = segment_csr(value, ptr, reduce=reduce)
             value = value[0] if isinstance(value, tuple) else value

torch_sparse/tensor.py

@@ -41,9 +41,10 @@ def from_edge_index(self, edge_index: torch.Tensor,
     @classmethod
     def from_dense(self, mat: torch.Tensor, has_value: bool = True):
         if mat.dim() > 2:
-            index = mat.abs().sum([i for i in range(2, mat.dim())]).nonzero()
+            index = mat.abs().sum([i for i in range(2, mat.dim())
+                                   ]).nonzero(as_tuple=False)
         else:
-            index = mat.nonzero()
+            index = mat.nonzero(as_tuple=False)
         index = index.t()
 
         row = index[0]