fix: update max_iter, related to 95cde5 clovaai#29

related to 95cde5 clovaai#29
stungkit · Aug 31, 2022 · 392ed80 · 392ed80
1 parent 8f2d515
commit 392ed80
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/lightning_module.py b/lightning_module.py
@@ -107,13 +107,13 @@ def configure_optimizers(self):
 
         max_iter = None
 
-        if self.config.get("max_epochs", None):
+        if int(self.config.get("max_epochs", -1)) > 0:
             assert len(self.config.train_batch_sizes) == 1, "Set max_epochs only if the number of datasets is 1"
             max_iter = (self.config.max_epochs * self.config.num_training_samples_per_epoch) / (
                 self.config.train_batch_sizes[0] * torch.cuda.device_count() * self.config.get("num_nodes", 1)
             )
 
-        if self.config.get("max_steps", None):
+        if int(self.config.get("max_steps", -1)) > 0:
             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