Skip to content

Commit

Permalink
fix bug when total_training_steps is not specified
Browse files Browse the repository at this point in the history
If `total_training_steps` is not specified, the default will be -1.
Thus checking whether it is > 0 is needed
  • Loading branch information
rayandrew committed Oct 16, 2024
1 parent f123cda commit d352b5d
Showing 1 changed file with 1 addition and 1 deletion.
2 changes: 1 addition & 1 deletion dlio_benchmark/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -257,7 +257,7 @@ def _train(self, epoch):
loader = self.framework.get_loader(dataset_type=DatasetType.TRAIN)
t0 = time()
for batch in dlp.iter(loader.next()):
if overall_step > max_steps or overall_step > self.total_training_steps:
if overall_step > max_steps or ((self.total_training_steps > 0) and (overall_step > self.total_training_steps)):
if self.args.my_rank == 0:
logging.info(f"{utcnow()} Maximum number of steps reached")
if (block_step != 1 and self.do_checkpoint) or (not self.do_checkpoint):
Expand Down

0 comments on commit d352b5d

Please sign in to comment.