thank you @wayi for this. I dont know why I didnt think to try thinking it could be inplace!
Just as a follow up quetsion, this wrapper basically does the same as follows?
        new_state_dict = collections.OrderedDict()
        for k, v in checkpoint['model_dict'].items():
            name = k.replace("module.", '') # remove `module.`
            new_state_dict[name] = v

Is that right? Are there any advantages/disadvantages of using either approaches?
Thank you again!