feat: Ensure that training resumes from latest checkpoint (#63)

* feat: Ensure that training resumes from latest checkpoint

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fix lint issues

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

* Fix import order

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

---------

Signed-off-by: Pranav Prashant Thombre <pthombre@nvidia.com>

Author: pthombre

dfm/src/automodel/recipes/train.py

@@ -306,9 +306,8 @@ def setup(self):
             start_epoch=int(self.start_epoch),
             num_epochs=int(self.num_epochs),
         )
-        # Optional resume only through config-defined restore_from
-        if self.restore_from:
-            self.load_checkpoint(restore_from=self.restore_from)
+
+        self.load_checkpoint(self.restore_from)
 
         if is_main_process():
             os.makedirs(self.checkpoint_config.checkpoint_dir, exist_ok=True)

examples/automodel/finetune/finetune.py

@@ -19,7 +19,7 @@
 from dfm.src.automodel.recipes.train import TrainWan21DiffusionRecipe
 
 
-def main(default_config_path="/opt/DFM/dfm/examples/Automodel/finetune/wan2_1_t2v_flow.yaml"):
+def main(default_config_path="examples/automodel/finetune/wan2_1_t2v_flow.yaml"):
     cfg = parse_args_and_load_config(default_config_path)
     recipe = TrainWan21DiffusionRecipe(cfg)
     recipe.setup()

examples/automodel/pretrain/pretrain.py

@@ -19,7 +19,7 @@
 from dfm.src.automodel.recipes.train import TrainWan21DiffusionRecipe
 
 
-def main(default_config_path="/opt/DFM/dfm/examples/Automodel/pretrain/wan2_1_t2v_flow.yaml"):
+def main(default_config_path="examples/automodel/pretrain/wan2_1_t2v_flow.yaml"):
     cfg = parse_args_and_load_config(default_config_path)
     recipe = TrainWan21DiffusionRecipe(cfg)
     recipe.setup()