Move save config before training start log (#478)

* Move save config before training start log

* Update tools/train.py

Co-authored-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

---------

Co-authored-by: Rustam Khadipash <16683750+hadipash@users.noreply.github.com>

Author: zhtmike

tools/train.py

@@ -195,6 +195,11 @@ def main(cfg):
         start_epoch=start_epoch,
     )
 
+    # save args used for training
+    if rank_id in [None, 0]:
+        with open(os.path.join(cfg.train.ckpt_save_dir, "args.yaml"), "w") as f:
+            yaml.safe_dump(cfg.to_dict(), stream=f, default_flow_style=False, sort_keys=False)
+
     # log
     num_devices = device_num if device_num is not None else 1
     global_batch_size = cfg.train.loader.batch_size * num_devices * gradient_accumulation_steps
@@ -229,12 +234,6 @@ def main(cfg):
         f"\nStart training... (The first epoch takes longer, please wait...)\n"
     )
 
-    # save args used for training
-    if rank_id in [None, 0]:
-        with open(os.path.join(cfg.train.ckpt_save_dir, "args.yaml"), "w") as f:
-            args_text = yaml.safe_dump(cfg.to_dict(), default_flow_style=False, sort_keys=False)
-            f.write(args_text)
-
     # training
     model = ms.Model(train_net)
     model.train(