Yes, for DataParallel, if you save by torch.save(model.state_dict()), it will save parameters on GPU 0.
But the parameters will be saved under model.module which cannot be loaded to non-DataParallel formats. That’s why I suggest the above code that makes saving/loading compatible with nn.Module format and nn.DataParallel format.
If you use DistributedDataParallel, you should only save parameters from local rank 0. Otherwise, each DDP process will try to save model on each GPU and overwrite each other.