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