unexpected GPU usage of cuda:0 in a multi-GPU env

Hi, I recently try the ViTDet project https://github.com/facebookresearch/detectron2/tree/main/projects/ViTDet and train the models via the training command with a checkpoint loaded:
```
../../tools/lazyconfig_train_net.py --config-file configs/path/to/config.py train.init_checkpoint=some/model
```
The training works, but I noticed that the usage of the first GPU cuda:0 is very high compared with the rest, say 38GB/40GB vs. 11GB/40GB. This limits the max batch size I can set and the validation process, even though there are still space in other GPUs.

I tested a bit and found out that the usage of cuda:0 becomes huge right after checkpoint load:
```
checkpointer.resume_or_load(cfg.train.init_checkpoint, resume=args.resume)
```

I have no idea how to solve this and I do not have access to edit my local env source code of detectron2.

Any help and suggestions will be greatly appreciated.