Cannot send models to gpu under MIG on A100

1

Right now there are three models, convnext_base, vgg16, and deit_base.
The first two models are from Pytorch.pretrained, and the last one is from HuggingFace.
The only change I made is I changed the output feature number of the last layer of VGG16 from 1000 to 10, since I am testing it on CIFAR-10.

With MIG on, and pass CUDA_VISIBLE_DEVICES=$UUID to the program, I can successfully send convnext_base and deit_base to the device. However, VGG16 failed, with error returned that

RuntimeError: Attempting to deserialize object on CUDA device 0 but torch.cuda.device_count() is 0. Please use torch.load with map_location to map your storages to an existing device.

Best
Max

2

Could you post the commands you’ve used to setup MIG on your A100, so that I could try to reproduce it, please?
Also, could you post the output of python -m torch.utils.collect_env?