search

rusty1s / pytorch_sparse

PyTorch Extension Library of Optimized Autograd Sparse Matrix Operations
MIT License
1.01k stars 147 forks source link

OOM for spmm but not matrix.mm #249

Closed justanhduc closed 2 years ago

justanhduc commented 2 years ago

Hi @rusty1s. Thanks for making this repo. I tried to replace Pytorch mm by spmm in my project but I got OOM. Is it already well-known or should the memory usage be similar somewhat?

rusty1s commented 2 years ago

You can try to use SparseTensor instead:

from torch_sparse import SparseTensor
sparse_mat = SparseTensor(row=..., col=..., value=..., size=(N, N))
out = sparse_mat @ dense_mat
justanhduc commented 2 years ago

Great it works! Thanks!