From 91c8d1ca1329644d75853fe8992d6aac03b2a992 Mon Sep 17 00:00:00 2001 From: zion <51308183+nousr@users.noreply.github.com> Date: Fri, 11 Nov 2022 12:15:13 -0800 Subject: [PATCH] bug fix cosine annealing optimizer in prior trainer (#262) --- dalle2_pytorch/trainer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/dalle2_pytorch/trainer.py b/dalle2_pytorch/trainer.py index a6b6bbc..b20b206 100644 --- a/dalle2_pytorch/trainer.py +++ b/dalle2_pytorch/trainer.py @@ -236,7 +236,7 @@ class DiffusionPriorTrainer(nn.Module): ) if exists(cosine_decay_max_steps): - self.scheduler = CosineAnnealingLR(optimizer, T_max = cosine_decay_max_steps) + self.scheduler = CosineAnnealingLR(self.optimizer, T_max = cosine_decay_max_steps) else: self.scheduler = LambdaLR(self.optimizer, lr_lambda = lambda _: 1.0)