Repository: pytorch
Issue: High occupation on GPU 0 when converting Tensor to multi GPU
Labels: module: performance,module: cuda,triaged

### 🐛 Describe the bug

I'm using distributed GPU training, and using `args.device = torch.device(f'cuda:{args.local_rank}' if torch.cuda.is_available() else 'cpu')` to set device. When I convert model to GPU with `model.to(args.device)`, it works well. However, when I convert data  to GPU during training with the following code:

```python
if isinstance(x, list):
    for x_i in range(len(x)):
        x[x_i] = x[x_i].to(args.device, non_blocking=True)
        label[x_i] = label[x_i].to(args.device, non_blocking=True)
    x = torch.cat(x, dim=0)
    label = torch.cat(label, dim=0)
else:
    x = x.to(args.device, non_blocking=True)  # .cuda()
    label = label.to(args.device, non_blocking=True)  # .cuda()
```

It makes `GPU 0` have much higher occupation than other GPUs, like the following picture, `3488` vs `1443`. And all the tensors are converted from CPU to the correct GPUs, which means the code assigned `args.device` correctly.

![image](https://user-images.githubusercontent.com/93592812/193007766-88dbc306-8393-4bd3-a3be-05894148d9af.png)

But when I set device with `torch.cuda.set_device(args.local_rank)`  at the beginning of my code and convert cpu tensor to gpu with `x.cuda(non_blocking=True)`, all the GPUs have the same memory-usage with `1443`

### Versions

PyTorch version: 1.8.1
Is debug build: False
CUDA used to build PyTorch: 10.2
ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64)
GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0
Clang version: Could not collect
CMake version: version 3.10.2
Libc version: glibc-2.27

Python version: 3.8.13 | packaged by conda-forge | (default, Mar 25 2022, 06:04:18)  [GCC 10.3.0] (64-bit runtime)
Python platform: Linux-4.19.126-1.jdcloud.x86_64-x86_64-with-glibc2.10
Is CUDA available: True
CUDA runtime version: 10.2.89
GPU models and configuration: 
GPU 0: Tesla P40
GPU 1: Tesla P40
GPU 2: Tesla P40
GPU 3: Tesla P40

Nvidia driver version: 470.57.02
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.3.1
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.3.1
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.22.4
[pip3] torch==1.8.1
[pip3] torchaudio==0.8.0a0+e4e171a
[pip3] torchvision==0.9.1
[conda] blas                      1.0                         mkl    conda-forge
[conda] cudatoolkit               10.2.89             h713d32c_10    conda-forge
[conda] ffmpeg                    4.3                  hf484d3e_0    pytorch
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           hc2b9512_224    defaults
[conda] numpy                     1.22.4           py38h99721a1_0    conda-forge
[conda] pytorch                   1.8.1           py3.8_cuda10.2_cudnn7.6.5_0    pytorch
[conda] torchaudio                0.8.1                      py38    pytorch
[conda] torchvision               0.9.1                py38_cu102    pytorch

cc @VitalyFedyunin @ngimel