Skip to content

GitLab

Unverified Commit 392ed80c authored by Geewook Kim's avatar Geewook Kim Committed by GitHub
Browse files

fix: update max_iter, related to 95cde5 #29 

related to 95cde5 #29
parent 8f2d5153
Show whitespace changes
Inline Side-by-side
Showing with 2 additions and 2 deletions
lightning_module.py
View file @ 392ed80c
...@@ -107,13 +107,13 @@ class DonutModelPLModule(pl.LightningModule): ...@@ -107,13 +107,13 @@ class DonutModelPLModule(pl.LightningModule):
max_iter = None 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" 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) / ( 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) 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 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 assert max_iter is not None
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment