ATen/SparseTensorUtils.h: No such file or directory

[EauDeData]

I've been running through the an issue when running AStartNet, raised from torchdrug.
I install the dependencies as follow, from a py3.10 fresh new conda environment

conda create -n py10 python==3.10; conda activate py10
pip install numpy==1.26.4
pip install torch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 --index-url https://download.pytorch.org/whl/cu124
pip install torch-scatter torch-sparse -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
pip install torch-cluster -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
pip install pyg-lib -f https://data.pyg.org/whl/torch-2.4.0+cu124.html
pip install torchdrug

pip install ogb
pip install easydict
pip install PyYAML
pip install easydict

However, after running the basic command CUDA_VISIBLE_DEVICES=0 python script/run.py -c config/transductive/fb15k237_astarnet.yaml --gpus 0

Torchdrug raises the following error message: fatal error: ATen/SparseTensorUtils.h: No such file or directory. Which is a sad but true; however, I don't know what else to try and install and play around multiple dependencies to solve it.

Is it similar in a way from what you did experience?
Thanks

Last callback of the error:

In file included from /home/amolina/miniconda3/lib/python3.11/site-packages/torchdrug-0.2.1-py3.11.egg/torchdrug/layers/functional/extension/rspmm.cu:6:
/home/amolina/miniconda3/lib/python3.11/site-packages/torchdrug-0.2.1-py3.11.egg/torchdrug/layers/functional/extension/rspmm.h:6:10: fatal error: ATen/SparseTensorUtils.h: No such file or directory
    6 | #include <ATen/SparseTensorUtils.h>
      |          ^~~~~~~~~~~~~~~~~~~~~~~~~~

[zchuang720]

I've checked the commits submitted in the past. In the new version of PyTorch, the SparseTensorUtils.h file has been moved to the ATen/native directory. You can try to modify the header file paths of rspmm.h and spmm.h under torchdrug/layers/functional/extension/:

replace
#include <ATen/SparseTensorUtils.h>
with
#include <ATen/native/SparseTensorUtils.h>.

I managed to make the code run successfully after making such modifications.