There is a "CUDA out of memory" error arising when we run the DBP15K example. The system hint that the error happens in the "__topk_\" function of the dgmc.py file. We find that this function needs to establish a matrix space with the shape of (1,n_s,n_t,d), where n and d are the node number and dimension respectively. This is too huge to create. When we modify the function to standard matrix multiplication, this error is repaired. Although I think this modification is equivalent to the original function, I don't know whether this violates your original code design.
That is indeed more memory-friendly. I changed it according to your suggestions. For DBP15K, it is advised to make use of keops though for even less memory consumption.
There is a "CUDA out of memory" error arising when we run the DBP15K example. The system hint that the error happens in the "__topk_\" function of the dgmc.py file. We find that this function needs to establish a matrix space with the shape of (1,n_s,n_t,d), where n and d are the node number and dimension respectively. This is too huge to create. When we modify the function to standard matrix multiplication, this error is repaired. Although I think this modification is equivalent to the original function, I don't know whether this violates your original code design.