def METHOD_NAME(self, state_dict, prefix):
    new_state_dict = dict()
    # remove unnecessary 'module.'
    for k, v in state_dict.items():
        if k.startswith(prefix):
            new_state_dict[k[len(prefix):]] = v
        else:
            new_state_dict[k] = v
    return new_state_dict