Adapting SWA for pytorch lightning.
Hello!
I'm currently trying to add SWA support to my lightning module, following this post https://pytorch.org/blog/stochastic-weight-averaging-in-pytorch/
And I have problems with loading averaged weights after I finish training.
I do it like this:
def on_train_end(self): #hook for SWA weights swapping
        if self.swa_enabled:
            print("Swapping model weights to SWA averages.")
            self.trainer.optimizers[0].swap_swa_sgd()
            print("Saving model with SWA weights to another checkpoint")
            self.save_checkpoint(os.path.join(self.trainer.checkpoint_callback.dirpath,f"swa_model_epoch={self.current_checkpoint}.ckpt"))
And bump into following problems:

It returns error message 'MyCoolModel' object has no attribute 'save_checkpoint'   on the line self.save_checkpoint....., instead of saving the model as stated in https://pytorch-lightning.readthedocs.io/en/latest/weights_loading.html
2)Even if I succeed with the 1st issue, I'm still required to update BN instances. As I believe, the best way to do this is to just run forward pass over the whole training data. So, my question is: is there a simple way to run 1 epoch though training data without writing my own loops?
Thanks in advance