Bugs for GPU based ot.sinkhorn, and very slow speed for GPU based ot.sinkhorn2

[HelloWorldLTY]

Hi, I met a bug when I intend to use GPU to calculate sinkhorn OT:

RuntimeError: CUDA error: device-side assert triggered
CUDA kernel errors might be asynchronously reported at some other API call,so the stacktrace below might be incorrect.
For debugging consider passing CUDA_LAUNCH_BLOCKING=1

Also I will meet this bug after running ot.sinkhorn2 for a while.

The errors come from my inside network codes, which is caused by choosing ot.sinkorn. If I choose ot.emd, then I will not meet such an error. Could you please help me? Thanks.

[rflamary]

Hello @HelloWorldLTY thank for the issue.

GPU bug are notoriously hard to debug so we need a wroking (or mor eprecisely not working) exmaple to reproduce it and understand what is happening. I know it is probably hidden inside a NN loss during traianing but could you try to give us more information and an example of when it fails?

Note that we tetsed the function on GPU and the computational gain of GPU is indeed null/limited for small problems as shwon in this page of the documentation
https://pythonot.github.io/gen_modules/ot.backend.html#module-ot.backend