[BUG] unimol+: RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:{0,1,2,3} and cpu!

### Describe the bug

The command listed in the README for 'Training PCQM4MV2' fails with `RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:$i and cpu!` for each GPU, at the beginning of training.

Output (appending `&> error.txt` to final command in reproduction): [error.txt](https://github.com/user-attachments/files/20773643/error.txt)

### Uni-Mol Version

unimol+

### Expected behavior

Training completes without issue.

### To Reproduce

Dockerfile:
> FROM dptechnology/unicore:latest-pytorch2.1.0-cuda12.1-rdma
RUN pip install rdkit==2022.09.3
RUN ldconfig && apt-get clean && apt-get autoremove && rm -rf /var/lib/apt/lists/* /tmp/*

Command in container after downloading data as stated in README:
`bash train_pcq.sh pcqm4m-v2 save_dir 32` (changed from default of 128 due to GPU memory limitations)

### Environment

Docker image with minor modifications as listed above (changing base image since README states PyTorch>2.0.0 and original Dockerfile uses 1.12.1, and removing the final `conda clean` due to conda being not present), running in Ubuntu 22.04.4 LTS, Docker CE 28.0.1, with four NVIDIA GPUs (4090 D, Driver 535.183.01, CUDA 12.2) exposed via `--runtime nvidia --gpus all` and extra parameters `--ipc=host --ulimit memlock=-1 --ulimit stack=67108864` as prompted on entering container.

`pip freeze` in container: [freeze.txt](https://github.com/user-attachments/files/20773648/freeze.txt)

### Additional Context

I tested that the bert example from uni-core (https://github.com/dptech-corp/Uni-Core/tree/main/examples/bert) runs without issue on all available GPUs on identical configurations, and that the error does not occur when I run the same container on another device with only one CUDA-capable GPU.
~~I will test how it goes when I expose only one GPU to the container at a later point.~~ Issue does not occur when I expose only one GPU to the container `(-e NVIDIA_VISIBLE_DEVICES=0 --gpus 1)`