Added cli argument for wandb session name

2026-04-08 22:35:09 +00:00 · 2024-01-03 11:52:38 +02:00
parent 1a36f9dc65
commit ab716302e4
2 changed files with 8 additions and 0 deletions
--- a/train_network.py
+++ b/train_network.py
@@ -684,6 +684,8 @@ class NetworkTrainer:

        if accelerator.is_main_process:
            init_kwargs = {}
+            if args.wandb_run_name:
+                init_kwargs['wandb'] = {'name': args.wandb_run_name}
            if args.log_tracker_config is not None:
                init_kwargs = toml.load(args.log_tracker_config)
            accelerator.init_trackers(