For me it’s still not fitting into GPU memory, so I guess I’ll need some other approach.
Anyway, I just spent the last day with another gotcha that’s far from obvious: I was calling this sparse matrix multiply in two parts of the graph (it’s all part of a much larger model). Now on toy examples I wasn’t getting proper convergence with this model, even when I used dense tensors! When I then replaced the call of that self-implemented function with a normal torch.mm(), it worked fine.
After lots of swearing and staring, I finally realized in the car what was happening: somehow only one of the two paths are used in the backward phase. After doing some extra things, it finally worked.
All of this is quite irrelevant, except that it shows me that this doesn’t seem to be the intended way to implement the sparse mm autograd. I’m not really finding enough documentation or simple examples, though.