Thanks for the update. I’ve rechecked your initial code and remembered that I’ve seen a similar issue before in pretrainedmodels and guess you might also be hitting this issue.
It seems the repository is breaking nn.DataParallel, so you could either use another repo (e.g. torchvision.models) or use DistributedDataParallel instead (I haven’t verified that it’s working with pretrainedmodels, but it might).