Swa model: Unexpected key(s) in state_dict: "n_averaged"

vision kornia
1

I have followed the swa tutorial of pytorch and saved the state dictionary of the AveragedModel as described here: PyTorch 1.6 now includes Stochastic Weight Averaging | PyTorch

Now, in my standalone inference script, I load the backbone (mobilenet in this case) and try to load the weights of the swa model to it. but I get:

 Unexpected key(s) in state_dict: "n_averaged".

I tried finding resources to load swa model state dict but no luck :frowning:

Any pointers to docs or tips on how to do this would be helpful!