[BUG] TensorDict.to moves all tensors to cuda:0 if index not specified

[ctongfei]

Describe the bug

When moving a TensorDict to a CUDA device whose index is not specified, tensordict moves it to cuda:0, where it should be moved to the current device as selected by torch.cuda.current_device() or torch.cuda.get_node_local_rank().

To Reproduce

import tensordict
import torch.cuda

torch.cuda.set_device(1)
# sets the current device index to be 1

assert torch.tensor(1).to(device=torch.device('cuda')).device.index == 1  
# passes assertion

assert TensorDict({'key': torch.tensor(1)}).to(device=torch.device('cuda')).device.index == 1  
# fails assertion! returned index == 0 instead of 1 

Reason and Possible fixes

The bug appears at

tensordict/tensordict/utils.py

Line 1393 in d1363eb

device = torch.device("cuda:0")

:

It should be changed to

    if device and device.type == "cuda" and device.index is None:
        device = torch.device("cuda", index=torch.cuda.current_device())

Checklist

• I have checked that there is no similar issue in the repo (required)
• I have read the documentation (required)
• I have provided a minimal working example to reproduce the bug (required)