Rooba
/
lightning
mirror of https://github.com/Lightning-AI/lightning.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

bug fix #10872 (#10965) 

Co-authored-by: louie.kim <louie.kim@kakaocorp.comlouie.kim@kakaocorp.com>
Co-authored-by: Jirka <jirka.borovec@seznam.cz>
This commit is contained in:
Biho-Kim 2022-02-09 23:15:49 +09:00 committed by GitHub
parent e3820da28a
commit 24de29974c
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
pytorch_lightning/callbacks/stochastic_weight_avg.py
View File

@ -197,7 +197,7 @@ class StochasticWeightAveraging(Callback):
self.n_averaged = torch.tensor(0, dtype=torch.long, device=pl_module.device)
if self.swa_start <= trainer.current_epoch <= self.swa_end:
self.update_parameters(self._average_model, pl_module, self.n_averaged, self.avg_fn)
self.update_parameters(self._average_model, pl_module, self.n_averaged, self._avg_fn)
# Note: No > here in case the callback is saved with the model and training continues
if trainer.current_epoch == self.swa_end + 1: