import torch

state_dict = torch.load("nnunetmonaibundle/model/Task09_Spleen_Bundle/models/FL_global_model.pt")

/tmp/ipykernel_3211/2956352158.py:1: FutureWarning: You are using `torch.load` with `weights_only=False` (the current default value), which uses the default pickle module implicitly. It is possible to construct malicious pickle data which will execute arbitrary code during unpickling (See https://github.com/pytorch/pytorch/blob/main/SECURITY.md#untrusted-models for more details). In a future release, the default value for `weights_only` will be flipped to `True`. This limits the functions that could be executed during unpickling. Arbitrary objects will no longer be allowed to be loaded via this mode unless they are explicitly allowlisted by the user via `torch.serialization.add_safe_globals`. We recommend you start setting `weights_only=True` for any use case where you don't have full control of the loaded file. Please open an issue on GitHub for any issues related to this experimental feature.
  state_dict = torch.load("nnunetmonaibundle/model/Task09_Spleen_Bundle/models/FL_global_model.pt")

new_state_dict = {}

from odict import odict
new_state_dict['network_weights'] = odict()

for weight in state_dict['model']:
    new_state_dict['network_weights'][weight[len("_orig_mod."):]] = state_dict['model'][weight]

new_state_dict['network_weights'].keys()

['encoder.stages.0.0.convs.0.conv.weight',
 'encoder.stages.0.0.convs.0.conv.bias',
 'encoder.stages.0.0.convs.0.norm.weight',
 'encoder.stages.0.0.convs.0.norm.bias',
 'encoder.stages.0.0.convs.0.all_modules.0.weight',
 'encoder.stages.0.0.convs.0.all_modules.0.bias',
 'encoder.stages.0.0.convs.0.all_modules.1.weight',
 'encoder.stages.0.0.convs.0.all_modules.1.bias',
 'encoder.stages.0.0.convs.1.conv.weight',
 'encoder.stages.0.0.convs.1.conv.bias',
 'encoder.stages.0.0.convs.1.norm.weight',
 'encoder.stages.0.0.convs.1.norm.bias',
 'encoder.stages.0.0.convs.1.all_modules.0.weight',
 'encoder.stages.0.0.convs.1.all_modules.0.bias',
 'encoder.stages.0.0.convs.1.all_modules.1.weight',
 'encoder.stages.0.0.convs.1.all_modules.1.bias',
 'encoder.stages.1.0.convs.0.conv.weight',
 'encoder.stages.1.0.convs.0.conv.bias',
 'encoder.stages.1.0.convs.0.norm.weight',
 'encoder.stages.1.0.convs.0.norm.bias',
 'encoder.stages.1.0.convs.0.all_modules.0.weight',
 'encoder.stages.1.0.convs.0.all_modules.0.bias',
 'encoder.stages.1.0.convs.0.all_modules.1.weight',
 'encoder.stages.1.0.convs.0.all_modules.1.bias',
 'encoder.stages.1.0.convs.1.conv.weight',
 'encoder.stages.1.0.convs.1.conv.bias',
 'encoder.stages.1.0.convs.1.norm.weight',
 'encoder.stages.1.0.convs.1.norm.bias',
 'encoder.stages.1.0.convs.1.all_modules.0.weight',
 'encoder.stages.1.0.convs.1.all_modules.0.bias',
 'encoder.stages.1.0.convs.1.all_modules.1.weight',
 'encoder.stages.1.0.convs.1.all_modules.1.bias',
 'encoder.stages.2.0.convs.0.conv.weight',
 'encoder.stages.2.0.convs.0.conv.bias',
 'encoder.stages.2.0.convs.0.norm.weight',
 'encoder.stages.2.0.convs.0.norm.bias',
 'encoder.stages.2.0.convs.0.all_modules.0.weight',
 'encoder.stages.2.0.convs.0.all_modules.0.bias',
 'encoder.stages.2.0.convs.0.all_modules.1.weight',
 'encoder.stages.2.0.convs.0.all_modules.1.bias',
 'encoder.stages.2.0.convs.1.conv.weight',
 'encoder.stages.2.0.convs.1.conv.bias',
 'encoder.stages.2.0.convs.1.norm.weight',
 'encoder.stages.2.0.convs.1.norm.bias',
 'encoder.stages.2.0.convs.1.all_modules.0.weight',
 'encoder.stages.2.0.convs.1.all_modules.0.bias',
 'encoder.stages.2.0.convs.1.all_modules.1.weight',
 'encoder.stages.2.0.convs.1.all_modules.1.bias',
 'encoder.stages.3.0.convs.0.conv.weight',
 'encoder.stages.3.0.convs.0.conv.bias',
 'encoder.stages.3.0.convs.0.norm.weight',
 'encoder.stages.3.0.convs.0.norm.bias',
 'encoder.stages.3.0.convs.0.all_modules.0.weight',
 'encoder.stages.3.0.convs.0.all_modules.0.bias',
 'encoder.stages.3.0.convs.0.all_modules.1.weight',
 'encoder.stages.3.0.convs.0.all_modules.1.bias',
 'encoder.stages.3.0.convs.1.conv.weight',
 'encoder.stages.3.0.convs.1.conv.bias',
 'encoder.stages.3.0.convs.1.norm.weight',
 'encoder.stages.3.0.convs.1.norm.bias',
 'encoder.stages.3.0.convs.1.all_modules.0.weight',
 'encoder.stages.3.0.convs.1.all_modules.0.bias',
 'encoder.stages.3.0.convs.1.all_modules.1.weight',
 'encoder.stages.3.0.convs.1.all_modules.1.bias',
 'encoder.stages.4.0.convs.0.conv.weight',
 'encoder.stages.4.0.convs.0.conv.bias',
 'encoder.stages.4.0.convs.0.norm.weight',
 'encoder.stages.4.0.convs.0.norm.bias',
 'encoder.stages.4.0.convs.0.all_modules.0.weight',
 'encoder.stages.4.0.convs.0.all_modules.0.bias',
 'encoder.stages.4.0.convs.0.all_modules.1.weight',
 'encoder.stages.4.0.convs.0.all_modules.1.bias',
 'encoder.stages.4.0.convs.1.conv.weight',
 'encoder.stages.4.0.convs.1.conv.bias',
 'encoder.stages.4.0.convs.1.norm.weight',
 'encoder.stages.4.0.convs.1.norm.bias',
 'encoder.stages.4.0.convs.1.all_modules.0.weight',
 'encoder.stages.4.0.convs.1.all_modules.0.bias',
 'encoder.stages.4.0.convs.1.all_modules.1.weight',
 'encoder.stages.4.0.convs.1.all_modules.1.bias',
 'encoder.stages.5.0.convs.0.conv.weight',
 'encoder.stages.5.0.convs.0.conv.bias',
 'encoder.stages.5.0.convs.0.norm.weight',
 'encoder.stages.5.0.convs.0.norm.bias',
 'encoder.stages.5.0.convs.0.all_modules.0.weight',
 'encoder.stages.5.0.convs.0.all_modules.0.bias',
 'encoder.stages.5.0.convs.0.all_modules.1.weight',
 'encoder.stages.5.0.convs.0.all_modules.1.bias',
 'encoder.stages.5.0.convs.1.conv.weight',
 'encoder.stages.5.0.convs.1.conv.bias',
 'encoder.stages.5.0.convs.1.norm.weight',
 'encoder.stages.5.0.convs.1.norm.bias',
 'encoder.stages.5.0.convs.1.all_modules.0.weight',
 'encoder.stages.5.0.convs.1.all_modules.0.bias',
 'encoder.stages.5.0.convs.1.all_modules.1.weight',
 'encoder.stages.5.0.convs.1.all_modules.1.bias',
 'decoder.encoder.stages.0.0.convs.0.conv.weight',
 'decoder.encoder.stages.0.0.convs.0.conv.bias',
 'decoder.encoder.stages.0.0.convs.0.norm.weight',
 'decoder.encoder.stages.0.0.convs.0.norm.bias',
 'decoder.encoder.stages.0.0.convs.0.all_modules.0.weight',
 'decoder.encoder.stages.0.0.convs.0.all_modules.0.bias',
 'decoder.encoder.stages.0.0.convs.0.all_modules.1.weight',
 'decoder.encoder.stages.0.0.convs.0.all_modules.1.bias',
 'decoder.encoder.stages.0.0.convs.1.conv.weight',
 'decoder.encoder.stages.0.0.convs.1.conv.bias',
 'decoder.encoder.stages.0.0.convs.1.norm.weight',
 'decoder.encoder.stages.0.0.convs.1.norm.bias',
 'decoder.encoder.stages.0.0.convs.1.all_modules.0.weight',
 'decoder.encoder.stages.0.0.convs.1.all_modules.0.bias',
 'decoder.encoder.stages.0.0.convs.1.all_modules.1.weight',
 'decoder.encoder.stages.0.0.convs.1.all_modules.1.bias',
 'decoder.encoder.stages.1.0.convs.0.conv.weight',
 'decoder.encoder.stages.1.0.convs.0.conv.bias',
 'decoder.encoder.stages.1.0.convs.0.norm.weight',
 'decoder.encoder.stages.1.0.convs.0.norm.bias',
 'decoder.encoder.stages.1.0.convs.0.all_modules.0.weight',
 'decoder.encoder.stages.1.0.convs.0.all_modules.0.bias',
 'decoder.encoder.stages.1.0.convs.0.all_modules.1.weight',
 'decoder.encoder.stages.1.0.convs.0.all_modules.1.bias',
 'decoder.encoder.stages.1.0.convs.1.conv.weight',
 'decoder.encoder.stages.1.0.convs.1.conv.bias',
 'decoder.encoder.stages.1.0.convs.1.norm.weight',
 'decoder.encoder.stages.1.0.convs.1.norm.bias',
 'decoder.encoder.stages.1.0.convs.1.all_modules.0.weight',
 'decoder.encoder.stages.1.0.convs.1.all_modules.0.bias',
 'decoder.encoder.stages.1.0.convs.1.all_modules.1.weight',
 'decoder.encoder.stages.1.0.convs.1.all_modules.1.bias',
 'decoder.encoder.stages.2.0.convs.0.conv.weight',
 'decoder.encoder.stages.2.0.convs.0.conv.bias',
 'decoder.encoder.stages.2.0.convs.0.norm.weight',
 'decoder.encoder.stages.2.0.convs.0.norm.bias',
 'decoder.encoder.stages.2.0.convs.0.all_modules.0.weight',
 'decoder.encoder.stages.2.0.convs.0.all_modules.0.bias',
 'decoder.encoder.stages.2.0.convs.0.all_modules.1.weight',
 'decoder.encoder.stages.2.0.convs.0.all_modules.1.bias',
 'decoder.encoder.stages.2.0.convs.1.conv.weight',
 'decoder.encoder.stages.2.0.convs.1.conv.bias',
 'decoder.encoder.stages.2.0.convs.1.norm.weight',
 'decoder.encoder.stages.2.0.convs.1.norm.bias',
 'decoder.encoder.stages.2.0.convs.1.all_modules.0.weight',
 'decoder.encoder.stages.2.0.convs.1.all_modules.0.bias',
 'decoder.encoder.stages.2.0.convs.1.all_modules.1.weight',
 'decoder.encoder.stages.2.0.convs.1.all_modules.1.bias',
 'decoder.encoder.stages.3.0.convs.0.conv.weight',
 'decoder.encoder.stages.3.0.convs.0.conv.bias',
 'decoder.encoder.stages.3.0.convs.0.norm.weight',
 'decoder.encoder.stages.3.0.convs.0.norm.bias',
 'decoder.encoder.stages.3.0.convs.0.all_modules.0.weight',
 'decoder.encoder.stages.3.0.convs.0.all_modules.0.bias',
 'decoder.encoder.stages.3.0.convs.0.all_modules.1.weight',
 'decoder.encoder.stages.3.0.convs.0.all_modules.1.bias',
 'decoder.encoder.stages.3.0.convs.1.conv.weight',
 'decoder.encoder.stages.3.0.convs.1.conv.bias',
 'decoder.encoder.stages.3.0.convs.1.norm.weight',
 'decoder.encoder.stages.3.0.convs.1.norm.bias',
 'decoder.encoder.stages.3.0.convs.1.all_modules.0.weight',
 'decoder.encoder.stages.3.0.convs.1.all_modules.0.bias',
 'decoder.encoder.stages.3.0.convs.1.all_modules.1.weight',
 'decoder.encoder.stages.3.0.convs.1.all_modules.1.bias',
 'decoder.encoder.stages.4.0.convs.0.conv.weight',
 'decoder.encoder.stages.4.0.convs.0.conv.bias',
 'decoder.encoder.stages.4.0.convs.0.norm.weight',
 'decoder.encoder.stages.4.0.convs.0.norm.bias',
 'decoder.encoder.stages.4.0.convs.0.all_modules.0.weight',
 'decoder.encoder.stages.4.0.convs.0.all_modules.0.bias',
 'decoder.encoder.stages.4.0.convs.0.all_modules.1.weight',
 'decoder.encoder.stages.4.0.convs.0.all_modules.1.bias',
 'decoder.encoder.stages.4.0.convs.1.conv.weight',
 'decoder.encoder.stages.4.0.convs.1.conv.bias',
 'decoder.encoder.stages.4.0.convs.1.norm.weight',
 'decoder.encoder.stages.4.0.convs.1.norm.bias',
 'decoder.encoder.stages.4.0.convs.1.all_modules.0.weight',
 'decoder.encoder.stages.4.0.convs.1.all_modules.0.bias',
 'decoder.encoder.stages.4.0.convs.1.all_modules.1.weight',
 'decoder.encoder.stages.4.0.convs.1.all_modules.1.bias',
 'decoder.encoder.stages.5.0.convs.0.conv.weight',
 'decoder.encoder.stages.5.0.convs.0.conv.bias',
 'decoder.encoder.stages.5.0.convs.0.norm.weight',
 'decoder.encoder.stages.5.0.convs.0.norm.bias',
 'decoder.encoder.stages.5.0.convs.0.all_modules.0.weight',
 'decoder.encoder.stages.5.0.convs.0.all_modules.0.bias',
 'decoder.encoder.stages.5.0.convs.0.all_modules.1.weight',
 'decoder.encoder.stages.5.0.convs.0.all_modules.1.bias',
 'decoder.encoder.stages.5.0.convs.1.conv.weight',
 'decoder.encoder.stages.5.0.convs.1.conv.bias',
 'decoder.encoder.stages.5.0.convs.1.norm.weight',
 'decoder.encoder.stages.5.0.convs.1.norm.bias',
 'decoder.encoder.stages.5.0.convs.1.all_modules.0.weight',
 'decoder.encoder.stages.5.0.convs.1.all_modules.0.bias',
 'decoder.encoder.stages.5.0.convs.1.all_modules.1.weight',
 'decoder.encoder.stages.5.0.convs.1.all_modules.1.bias',
 'decoder.stages.0.convs.0.conv.weight',
 'decoder.stages.0.convs.0.conv.bias',
 'decoder.stages.0.convs.0.norm.weight',
 'decoder.stages.0.convs.0.norm.bias',
 'decoder.stages.0.convs.0.all_modules.0.weight',
 'decoder.stages.0.convs.0.all_modules.0.bias',
 'decoder.stages.0.convs.0.all_modules.1.weight',
 'decoder.stages.0.convs.0.all_modules.1.bias',
 'decoder.stages.0.convs.1.conv.weight',
 'decoder.stages.0.convs.1.conv.bias',
 'decoder.stages.0.convs.1.norm.weight',
 'decoder.stages.0.convs.1.norm.bias',
 'decoder.stages.0.convs.1.all_modules.0.weight',
 'decoder.stages.0.convs.1.all_modules.0.bias',
 'decoder.stages.0.convs.1.all_modules.1.weight',
 'decoder.stages.0.convs.1.all_modules.1.bias',
 'decoder.stages.1.convs.0.conv.weight',
 'decoder.stages.1.convs.0.conv.bias',
 'decoder.stages.1.convs.0.norm.weight',
 'decoder.stages.1.convs.0.norm.bias',
 'decoder.stages.1.convs.0.all_modules.0.weight',
 'decoder.stages.1.convs.0.all_modules.0.bias',
 'decoder.stages.1.convs.0.all_modules.1.weight',
 'decoder.stages.1.convs.0.all_modules.1.bias',
 'decoder.stages.1.convs.1.conv.weight',
 'decoder.stages.1.convs.1.conv.bias',
 'decoder.stages.1.convs.1.norm.weight',
 'decoder.stages.1.convs.1.norm.bias',
 'decoder.stages.1.convs.1.all_modules.0.weight',
 'decoder.stages.1.convs.1.all_modules.0.bias',
 'decoder.stages.1.convs.1.all_modules.1.weight',
 'decoder.stages.1.convs.1.all_modules.1.bias',
 'decoder.stages.2.convs.0.conv.weight',
 'decoder.stages.2.convs.0.conv.bias',
 'decoder.stages.2.convs.0.norm.weight',
 'decoder.stages.2.convs.0.norm.bias',
 'decoder.stages.2.convs.0.all_modules.0.weight',
 'decoder.stages.2.convs.0.all_modules.0.bias',
 'decoder.stages.2.convs.0.all_modules.1.weight',
 'decoder.stages.2.convs.0.all_modules.1.bias',
 'decoder.stages.2.convs.1.conv.weight',
 'decoder.stages.2.convs.1.conv.bias',
 'decoder.stages.2.convs.1.norm.weight',
 'decoder.stages.2.convs.1.norm.bias',
 'decoder.stages.2.convs.1.all_modules.0.weight',
 'decoder.stages.2.convs.1.all_modules.0.bias',
 'decoder.stages.2.convs.1.all_modules.1.weight',
 'decoder.stages.2.convs.1.all_modules.1.bias',
 'decoder.stages.3.convs.0.conv.weight',
 'decoder.stages.3.convs.0.conv.bias',
 'decoder.stages.3.convs.0.norm.weight',
 'decoder.stages.3.convs.0.norm.bias',
 'decoder.stages.3.convs.0.all_modules.0.weight',
 'decoder.stages.3.convs.0.all_modules.0.bias',
 'decoder.stages.3.convs.0.all_modules.1.weight',
 'decoder.stages.3.convs.0.all_modules.1.bias',
 'decoder.stages.3.convs.1.conv.weight',
 'decoder.stages.3.convs.1.conv.bias',
 'decoder.stages.3.convs.1.norm.weight',
 'decoder.stages.3.convs.1.norm.bias',
 'decoder.stages.3.convs.1.all_modules.0.weight',
 'decoder.stages.3.convs.1.all_modules.0.bias',
 'decoder.stages.3.convs.1.all_modules.1.weight',
 'decoder.stages.3.convs.1.all_modules.1.bias',
 'decoder.stages.4.convs.0.conv.weight',
 'decoder.stages.4.convs.0.conv.bias',
 'decoder.stages.4.convs.0.norm.weight',
 'decoder.stages.4.convs.0.norm.bias',
 'decoder.stages.4.convs.0.all_modules.0.weight',
 'decoder.stages.4.convs.0.all_modules.0.bias',
 'decoder.stages.4.convs.0.all_modules.1.weight',
 'decoder.stages.4.convs.0.all_modules.1.bias',
 'decoder.stages.4.convs.1.conv.weight',
 'decoder.stages.4.convs.1.conv.bias',
 'decoder.stages.4.convs.1.norm.weight',
 'decoder.stages.4.convs.1.norm.bias',
 'decoder.stages.4.convs.1.all_modules.0.weight',
 'decoder.stages.4.convs.1.all_modules.0.bias',
 'decoder.stages.4.convs.1.all_modules.1.weight',
 'decoder.stages.4.convs.1.all_modules.1.bias',
 'decoder.transpconvs.0.weight',
 'decoder.transpconvs.0.bias',
 'decoder.transpconvs.1.weight',
 'decoder.transpconvs.1.bias',
 'decoder.transpconvs.2.weight',
 'decoder.transpconvs.2.bias',
 'decoder.transpconvs.3.weight',
 'decoder.transpconvs.3.bias',
 'decoder.transpconvs.4.weight',
 'decoder.transpconvs.4.bias',
 'decoder.seg_layers.0.weight',
 'decoder.seg_layers.0.bias',
 'decoder.seg_layers.1.weight',
 'decoder.seg_layers.1.bias',
 'decoder.seg_layers.2.weight',
 'decoder.seg_layers.2.bias',
 'decoder.seg_layers.3.weight',
 'decoder.seg_layers.3.bias',
 'decoder.seg_layers.4.weight',
 'decoder.seg_layers.4.bias']

torch.save(new_state_dict, "nnunetmonaibundle/model/Task09_Spleen_Bundle/models/FL_global_model_FINAL.pt")