@tom Hi, Thomas, thanks a lot for your response.
Actually I am having trouble with saving not the loading, namely:
torch.save(model.module.state_dict(), PATH)

Do you happen to know if this function only apply to the latest pytorch version? Cause I am using version ‘1.5.0a0+8f84ded’ and got the following error when it comes to saving checkpoint:
Traceback (most recent call last):
  File "train.py", line 227, in <module>
    main()
  File "train.py", line 224, in main
    run(config)
  File "train.py", line 206, in run
    state_dict, config, experiment_name)
  File "/home/BIGGAN/train_fns.py", line 101, in save_and_sample
    experiment_name, None, G_ema if config['ema'] else None)
  File "/home/BIGGAN/utils.py", line 721, in save_weights
    torch.save(G.module.state_dict(), 
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 594, in __getattr__
    type(self).__name__, name))
AttributeError: 'Generator' object has no attribute 'module'

Or do I need to do some extra configuration when setting up the models if I wanna use torch.save(model.module.state_dict(), PATH) to save torch.nn.DataParallel Models? Thank you!