clean up the_cauldron_transform args

pytorch · Nov 26, 2024 · cb67f43 · cb67f43
1 parent 7057aa9
commit cb67f43
Showing 1 changed file with 6 additions and 6 deletions.
diff --git a/recipes/lora_finetune_distributed_multi_dataset.py b/recipes/lora_finetune_distributed_multi_dataset.py
@@ -948,12 +948,12 @@ def train(self) -> None:
                     # will include multiple forward / backward passes if gradient accumulation > 1
                     self._profiler.step()
 
-                if self._is_rank_zero:
-                    total_twfb_pct = round(epoch_total_dl_time / epoch_step_time, 5)
-                    log.info(
-                        f"End of epoch {self.epochs_run}! "
-                        f"{total_twfb_pct=}, {epoch_step_time=}, {epoch_dl_time=}, {epoch_total_dl_time=}",
-                    )
+            if self._is_rank_zero:
+                total_twfb_pct = round(epoch_total_dl_time / epoch_step_time, 5)
+                log.info(
+                    f"End of epoch {self.epochs_run}! "
+                    f"{total_twfb_pct=}, {epoch_step_time=}, {epoch_dl_time=}, {epoch_total_dl_time=}",
+                )
             self.epochs_run += 1
             self.save_checkpoint(epoch=curr_epoch)