feat: support wandb logging

2026-04-08 22:35:09 +00:00 · 2023-04-20 01:41:12 +09:00
parent 334589af4e
commit 27ffd9fe3d
6 changed files with 33 additions and 7 deletions
--- a/train_db.py
+++ b/train_db.py
@@ -231,7 +231,7 @@ def train(args):
    )

    if accelerator.is_main_process:
-        accelerator.init_trackers("dreambooth")
+        accelerator.init_trackers("dreambooth" if args.log_tracker_name is None else args.log_tracker_name)

    loss_list = []
    loss_total = 0.0