[MRG] Gromov_Wasserstein2 not performing backward properly on GPU#352
Conversation
|
Thanks for the quick PR! One thing I wonder though is that how to check if the build on Circle CI does run the test with GPU? https://github.com/PythonOT/POT/runs/5389073677?check_suite_focus=true Because if you do this on the test file if torch.cuda.is_available():
devices.append(torch.device("cuda"))This line might not be hit at all during running the test file |
|
No you're right, tests performed here do not have access to a GPU. But I ran it on my machine which is equipped with V100s and the test passes. As always with GPU tests, they should be run locally to detect any issue. |
The tests passed on my local run as well, thanks. |

Types of changes
The backpropagation is not working in gromov_wasserstein2 if the given tensors are located on a GPU. This is due to the fact that part of the computation is performed with numpy and the device was forgotten when casting back to torch.
Motivation and context / Related issue
Resolves #351
How has this been tested (if it applies)
PR checklist