ot.solve_sample_batch devices error

[khoatranrb]

Describe the bug

I applied ot.solve_sample_batch but received an error RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!

To Reproduce

Code sample

import ot
import torch

x = torch.zeros((24, 100, 1))
y = torch.zeros((24, 100, 1))

ot.solve_sample_batch(x.cuda(), y.cuda(), reg=1)

Screenshots

Environment (please complete the following information):

I got that error from both Colab and Linux.

[rflamary]

Thanks, I think this can come from the reg value, could you try to make is a tensor on cuda also? We will look into it

[khoatranrb]

Thanks, I think this can come from the reg value, could you try to make is a tensor on cuda also? We will look into it

It comes from loga and logb from bregman_log_projection_batch of ot/batch/_utils.py. Hope u update it soon.

[rflamary]

In fact I think this comes from a and b set as None and probably initialized on the wrong device. If you give them it should work.

it comes from
https://github.com/PythonOT/POT/blob/master/ot/batch/_linear.py#L313

@KrzakalaPaul can you please do a PR and use type_asfor the initialization?

[rflamary]

It should be fixed on master branch now.