About the pytorch version

[siddharth9820]

Hi,
We are interested in reproducing your results on our local multi-gpu clusters. We wanted to understand if the version of pytorch still needs to be 1.3.1 or can we use a more recent version (2.0+)?

[alokpathy]

Hi Siddharth! Thanks for being patient with a response and for reaching out. I've enjoyed reading a lot of your work like AxoNN.

You can use a newer version of pytorch (e.g. 2.0+), though you'll need a newer version of the CUDA extension. I think if you replace sparse_coo_tensor.cu with this file to account for the newer CUDA version (tested on pytorch 2.1 and CUDA 12.2), things should run.

https://gist.github.com/alokpathy/8727cff462ac763a4c7d970a79a2166a