faster transpose

torch_sparse/transpose.py

@@ -29,11 +29,15 @@ def transpose(index, value, m, n, coalesced=True):
 
 
 def t(src):
-    (row, col), value = src.coo()
+    index, value = src.coo()
     csr2csc = src.storage.csr2csc
 
+    new_index = torch.empty_like(index)
+    torch.index_select(index[1], 0, csr2csc, out=new_index[0])
+    torch.index_select(index[0], 0, csr2csc, out=new_index[1])
+
     storage = src.storage.__class__(
-        index=torch.stack([col, row], dim=0)[:, csr2csc],
+        index=new_index,
         value=value[csr2csc] if src.has_value() else None,
         sparse_size=src.sparse_size()[::-1],
         rowcount=src.storage._colcount,