From 24de29974c24d6ccbab3c4d2377d5441a379e797 Mon Sep 17 00:00:00 2001 From: Biho-Kim <22164993+qqueing@users.noreply.github.com> Date: Wed, 9 Feb 2022 23:15:49 +0900 Subject: [PATCH] bug fix #10872 (#10965) Co-authored-by: louie.kim Co-authored-by: Jirka --- pytorch_lightning/callbacks/stochastic_weight_avg.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/pytorch_lightning/callbacks/stochastic_weight_avg.py b/pytorch_lightning/callbacks/stochastic_weight_avg.py index 42dd67b724..d18c9dcffd 100644 --- a/pytorch_lightning/callbacks/stochastic_weight_avg.py +++ b/pytorch_lightning/callbacks/stochastic_weight_avg.py @@ -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: