fix: update max_iter, related to 95cde5 #29

related to 95cde5 #29

lightning_module.py

@@ -107,13 +107,13 @@ class DonutModelPLModule(pl.LightningModule):
 
         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