he docs suggest to use:
To let a non-DDP model load a state dict from a DDP model, consume_prefix_in_state_dict_if_present() needs to be applied to strip the prefix “module.” in the DDP state dict before loading.

Source: DistributedDataParallel — PyTorch 1.9.0 documentation 7
however., this line:
model_state_dict=torch.nn.modules.utils.consume_prefix_in_state_dict_if_present(checkpoint['model_dict'],prefix='module.')

gives back None 
I do something ugly like this at the moment:
        new_state_dict = collections.OrderedDict()
        for k, v in checkpoint['model_dict'].items():
            name = k.replace("module.", '') # remove `module.`
            new_state_dict[name] = v

but I would prefer to use the consume_prefix_in_state_dict_if_present.
Can someone elucidate the correct usage of this please? Obviously, I am not getting it!