Sorry, @VoVAllen I found that I still have this issue: Error when train DistSAGE with OGB-Product
I wonder what is the dgl version in: [Bugfix] Fix CUDA 11.1 crashing when number of edges is larger than number of node pairs by BarclayII · Pull Request #3265 · dmlc/dgl · GitHub
Maybe I can try if this version would work.