fix: max_iter with config.max_steps

clovaai · Jul 29, 2022 · 4130973 · 4130973
1 parent 537cc20
commit 4130973
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/lightning_module.py b/lightning_module.py
@@ -114,7 +114,7 @@ def configure_optimizers(self):
             )
 
         if self.config.get("max_steps", None):
-            max_iter = min(self.config.max_steps, max_iter)
+            max_iter = min(self.config.max_steps, max_iter) if max_iter is not None else self.config.max_steps
 
         assert max_iter is not None
         optimizer = torch.optim.Adam(self.parameters(), lr=self.config.lr)