SWA callback with torch.save state_dict()

[maxmatical]

i'm trying to experiment with SWA callback when training pytorch models using the lightning trainer. one problem i'm running into is my lightningmodule is structured like

class MyModule(pl.LightningModule):
    def __init__(self, pytorch_model: nn.Module):
        self.pytorch_model = pytorch_model

    def forward(x):
        return self.pytorch_model(x)

...

my_module = MyModule(...)
trainer = Trainer(callbacks=[StochasticWeightAveraging(swa_epoch_start=0.5)])
trainer.fit(my_module)

torch.save(my_module.pytorch_model.state_dict(), save_path)

i tried training with and without swa callback (using the same seed for 10 epochs so i know it's averaging models). but after loading the pytorch model and running evaluation, i'm getting the same metric for the pytorch model when training with/without SWA

does anyone know why this is the case?

edit: solved the issue. turns out the trainer was automatically loading the best module checkpoint, which doesn't contain the averaged weights

[rohitgavval]

Could you please show how you were able to retrieve the model with the averaged weights, if you were able to? I can't seem to figure that out