diff --git a/CHANGELOG.md b/CHANGELOG.md index 0309e1e9ea..4490741efc 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -26,7 +26,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0 ### Fixed -- Fix doctest fails with ImportError: cannot import name 'Env' from 'gym' ([#751](https://github.com/PyTorchLightning/lightning-bolts/pull/751)) +- Fixed doctest fails with ImportError: cannot import name 'Env' from 'gym' ([#751](https://github.com/PyTorchLightning/lightning-bolts/pull/751)) + +- Fixed MoCo v2 missing Cosine Annealing learning scheduler ([#757](https://github.com/PyTorchLightning/lightning-bolts/pull/757)) ## [0.4.0] - 2021-09-09 diff --git a/pl_bolts/models/self_supervised/moco/moco2_module.py b/pl_bolts/models/self_supervised/moco/moco2_module.py index fd5bf4039a..5ab20d0d0f 100644 --- a/pl_bolts/models/self_supervised/moco/moco2_module.py +++ b/pl_bolts/models/self_supervised/moco/moco2_module.py @@ -308,7 +308,11 @@ def configure_optimizers(self): momentum=self.hparams.momentum, weight_decay=self.hparams.weight_decay, ) - return optimizer + scheduler = torch.optim.lr_scheduler.CosineAnnealingLR( + optimizer, + self.trainer.max_epochs, + ) + return [optimizer], [scheduler] @staticmethod def add_model_specific_args(parent_parser):