[BUG] Error loading Global best model with testing score script for Prostate3D example

[umgpy]

Describe the bug
When running testing score script from the example which has been modified for our data and task of binary classification, the model loading fails due to missing keys in the state dict.

To Reproduce
Steps to reproduce the behavior:

1. Go to '...'
2. Click on '....'
3. Scroll down to '....'
4. See error

Expected behavior
The best global model should load and run.

Screenshots
If applicable, add screenshots to help explain your problem.

Desktop (please complete the following information):

• OS: ubuntu 18.04
• Python Version: 3.9.12
• NVFlare Version:2.3.3
• Monai_flare: 0.2.3

Additional context
Error -

RuntimeError: Error(s) in loading state_dict for UNet:
	Missing key(s) in state_dict: "model.0.conv.unit0.conv.weight", "model.0.conv.unit0.conv.bias", "model.0.conv.unit0.adn.N.weight", "model.0.conv.unit0.adn.N.bias", "model.0.conv.unit0.adn.N.running_mean", "model.0.conv.unit0.adn.N.running_var", "model.0.conv.unit0.adn.A.weight", "model.0.conv.unit1.conv.weight", "model.0.conv.unit1.conv.bias", "model.0.conv.unit1.adn.N.weight", "model.0.conv.unit1.adn.N.bias", "model.0.conv.unit1.adn.N.running_mean", "model.0.conv.unit1.adn.N.running_var", "model.0.conv.unit1.adn.A.weight", "model.0.residual.weight", "model.0.residual.bias", "model.1.submodule.0.conv.unit0.conv.weight", "model.1.submodule.0.conv.unit0.conv.bias", "model.1.submodule.0.conv.unit0.adn.N.weight", "model.1.submodule.0.conv.unit0.adn.N.bias", "model.1.submodule.0.conv.unit0.adn.N.running_mean", "model.1.submodule.0.conv.unit0.adn.N.running_var", "model.1.submodule.0.conv.unit0.adn.A.weight", "model.1.submodule.0.conv.unit1.conv.weight", "model.1.submodule.0.conv.unit1.conv.bias", "model.1.submodule.0.conv.unit1.adn.N.weight", "model.1.submodule.0.conv.unit1.adn.N.bias", "model.1.submodule.0.conv.unit1.adn.N.running_mean", "model.1.submodule.0.conv.unit1.adn.N.running_var", "model.1.submodule.0.conv.unit1.adn.A.weight", "model.1.submodule.0.residual.weight", "model.1.submodule.0.residual.bias", "model.1.submodule.1.submodule.0.conv.unit0.conv.weight", "model.1.submodule.1.submodule.0.conv.unit0.conv.bias", "model.1.submodule.1.submodule.0.conv.unit0.adn.N.weight", "model.1.submodule.1.submodule.0.conv.unit0.adn.N.bias", "model.1.submodule.1.submodule.0.conv.unit0.adn.N.running_mean", "model.1.submodule.1.submodule.0.conv.unit0.adn.N.running_var", "model.1.submodule.1.submodule.0.conv.unit0.adn.A.weight", "model.1.submodule.1.submodule.0.conv.unit1.conv.weight", "model.1.submodule.1.submodule.0.conv.unit1.conv.bias", "model.1.submodule.1.submodule.0.conv.unit1.adn.N.weight", "model.1.submodule.1.submodule.0.conv.unit1.adn.N.bias", "model.1.submodule.1.submodule.0.conv.unit1.adn.N.running_mean", "model.1.submodule.1.submodule.0.conv.unit1.adn.N.running_var", "model.1.submodule.1.submodule.0.conv.unit1.adn.A.weight", "model.1.submodule.1.submodule.0.residual.weight", "model.1.submodule.1.submodule.0.residual.bias", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.weight", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.conv.bias", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.adn.N.weight", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.adn.N.bias", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.adn.N.running_mean", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.adn.N.running_var", "model.1.submodule.1.submodule.1.submodule.0.conv.unit0.adn.A.weight", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.weight", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.conv.bias", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.adn.N.weight", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.adn.N.bias", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.adn.N.running_mean", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.adn.N.running_var", "model.1.submodule.1.submodule.1.submodule.0.conv.unit1.adn.A.weight", "model.1.submodule.1.submodule.1.submodule.0.residual.weight", "model.1.submodule.1.submodule.1.submodule.0.residual.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.conv.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.adn.N.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.adn.N.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.adn.N.running_mean", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.adn.N.running_var", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit0.adn.A.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.conv.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.adn.N.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.adn.N.bias", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.adn.N.running_mean", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.adn.N.running_var", "model.1.submodule.1.submodule.1.submodule.1.submodule.conv.unit1.adn.A.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.residual.weight", "model.1.submodule.1.submodule.1.submodule.1.submodule.residual.bias", "model.1.submodule.1.submodule.1.submodule.2.0.conv.weight", "model.1.submodule.1.submodule.1.submodule.2.0.conv.bias", "model.1.submodule.1.submodule.1.submodule.2.0.adn.N.weight", "model.1.submodule.1.submodule.1.submodule.2.0.adn.N.bias", "model.1.submodule.1.submodule.1.submodule.2.0.adn.N.running_mean", "model.1.submodule.1.submodule.1.submodule.2.0.adn.N.running_var", "model.1.submodule.1.submodule.1.submodule.2.0.adn.A.weight", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.weight", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.conv.bias", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.adn.N.weight", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.adn.N.bias", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.adn.N.running_mean", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.adn.N.running_var", "model.1.submodule.1.submodule.1.submodule.2.1.conv.unit0.adn.A.weight", "model.1.submodule.1.submodule.2.0.conv.weight", "model.1.submodule.1.submodule.2.0.conv.bias", "model.1.submodule.1.submodule.2.0.adn.N.weight", "model.1.submodule.1.submodule.2.0.adn.N.bias", "model.1.submodule.1.submodule.2.0.adn.N.running_mean", "model.1.submodule.1.submodule.2.0.adn.N.running_var", "model.1.submodule.1.submodule.2.0.adn.A.weight", "model.1.submodule.1.submodule.2.1.conv.unit0.conv.weight", "model.1.submodule.1.submodule.2.1.conv.unit0.conv.bias", "model.1.submodule.1.submodule.2.1.conv.unit0.adn.N.weight", "model.1.submodule.1.submodule.2.1.conv.unit0.adn.N.bias", "model.1.submodule.1.submodule.2.1.conv.unit0.adn.N.running_mean", "model.1.submodule.1.submodule.2.1.conv.unit0.adn.N.running_var", "model.1.submodule.1.submodule.2.1.conv.unit0.adn.A.weight", "model.1.submodule.2.0.conv.weight", "model.1.submodule.2.0.conv.bias", "model.1.submodule.2.0.adn.N.weight", "model.1.submodule.2.0.adn.N.bias", "model.1.submodule.2.0.adn.N.running_mean", "model.1.submodule.2.0.adn.N.running_var", "model.1.submodule.2.0.adn.A.weight", "model.1.submodule.2.1.conv.unit0.conv.weight", "model.1.submodule.2.1.conv.unit0.conv.bias", "model.1.submodule.2.1.conv.unit0.adn.N.weight", "model.1.submodule.2.1.conv.unit0.adn.N.bias", "model.1.submodule.2.1.conv.unit0.adn.N.running_mean", "model.1.submodule.2.1.conv.unit0.adn.N.running_var", "model.1.submodule.2.1.conv.unit0.adn.A.weight", "model.2.0.conv.weight", "model.2.0.conv.bias", "model.2.0.adn.N.weight", "model.2.0.adn.N.bias", "model.2.0.adn.N.running_mean", "model.2.0.adn.N.running_var", "model.2.0.adn.A.weight", "model.2.1.conv.unit0.conv.weight", "model.2.1.conv.unit0.conv.bias". 
	Unexpected key(s) in state_dict: "meta_props", "train_conf". 

Model definition in SupervisedProstateLearner-

 self.model = UNet(
            spatial_dims=3,
            in_channels=1,
            out_channels=2,
            channels=(16, 32, 64, 128, 256),
            strides=(2, 2, 2, 2),
            num_res_units=2,
        ).to(self.device)

Model definition in testing script

model = UNet(
        spatial_dims=3,
        in_channels=1,
        out_channels=2,
        channels=(16, 32, 64, 128, 256),
        strides=(2, 2, 2, 2),
        num_res_units=2,
    ).to(device)
    model_weights = torch.load(model_path)
    model_weights = model_weights["model"]
    model.load_state_dict(model_weights)

[YuanTingHsieh]

@umgpy your NVFlare version is 2.3.3, so you want to try examples on branch 2.3 here: https://github.com/NVIDIA/NVFlare/tree/2.3/examples/advanced/prostate

[one2clouds]

Well one thing you could try checking is the shape of Unet architecture you defined while training is same while you define on testing. That solved my problem...