<jupyter_start><jupyter_text>Change state dicts' keys for `DistributedDataParallel` A model that is wrapped with `DistributedDataParallel` expects its state dict's keys to start with `'module.'`, followed by the usual keys when `DistributedDataParellel` is not used.  In order to load the pretrained models that come with the original ESRGAN-pytorch, which are not wrapped with `DistributedDataParallel`, their state dicts' keys are prepended with `'module.'` beforehand so that they can be loaded into models wrapped with `DistributedDataParallel`.  This is done once.  The resulting state dicts are saved to disk for future use.<jupyter_code>import torch
device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
models = ('psnr', 'gan')

for m in models:
    g_sdict = torch.load(f'../parameters/{m}.pth', map_location=device)
    g_sdict = {f'module.{k}':v for k, v in g_sdict.items()}
    save_dict = {
        'epoch':None, 
        'unit_scheduler_step':None, 
        'history':None,
        'g_state_dict':g_sdict, 
        'd_state_dict':None, 
        'opt_g_state_dict':None, 
        'opt_d_state_dict':None, 
        'amp':None, 
        'args':None}
    torch.save(save_dict, f'../parameters/{m}_based.pth')
! du -hs ../parameters/*<jupyter_output>64M	../parameters/gan.pth
 64M	../parameters/gan_based.pth
 64M	../parameters/psnr.pth
 64M	../parameters/psnr_based.pth
{'fraction_non_alphanumeric': 0.11928571428571429, 'fraction_numerical': 0.0064285714285714285, 'mean_word_length': 3.781569965870307, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '452836', 'license': 'MIT', 'n_tokens_mistral': 473, 'n_tokens_neox': 458, 'n_words': 132, 'path': 'nbs/01_pretrained_models_formats.ipynb', 'repo_name': 'qAp/ESRGAN-pytorch', 'size': '1400'}