Ok, I didn’t know that functions torch.ones_like would automatically put things on same device. Thanks!
But I think that we’re still missing something
with torch.cuda(device):
   whatever happens here make them automatically on cuda device defined

or something like
@torch.cuda(device)
def some_func():
  send whatever is in here to cuda device

Unless there exists something for such cases and I’ve completely missed it?