coalesce call

ffec0a56 · rusty1s · 5788c855 · ffec0a56 · 5788c855
Commit ffec0a56 authored Jul 30, 2018 by rusty1s
Hide whitespace changes
Inline Side-by-side

Showing with 3 additions and 7 deletions

cuda/matmul_cuda.cu cuda/matmul_cuda.cu +3 -0

torch_sparse/sparse.py torch_sparse/sparse.py +0 -7

No files found.
--- a/cuda/matmul_cuda.cu
+++ b/cuda/matmul_cuda.cu
@@ -28,6 +28,9 @@ static void init_cusparse() {
 }

 std::tuple<at::Tensor, at::Tensor> spspmm_cuda(at::Tensor A, at::Tensor B) {
+  A = A.coalesce();
+  B = B.coalesce();
+
  init_cusparse();

  auto m = A.size(0);

--- a/torch_sparse/sparse.py
+++ b/torch_sparse/sparse.py
-import torch
-
-
-def SparseTensor(index, value, size):
-    t = torch.cuda if value.is_cuda else torch
-    SparseTensor = getattr(t.sparse, value.type().split('.')[-1])
-    return SparseTensor(index, value, size)