From da979018316685d53aa75e94cef80393ab1a53b2 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <2225664821@qq.com>
Date: Wed, 23 Aug 2023 18:19:45 +0800
Subject: [PATCH] Update train_ms.py

---
 train_ms.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/train_ms.py b/train_ms.py
index a2f6ab0..d792833 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -8,6 +8,7 @@ from torch import nn, optim
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
 #from torch.utils.tensorboard import SummaryWriter
+import wandb
 from wandb.tensorboard import WandbSummaryWriter as SummaryWriter
 import torch.multiprocessing as mp
 import torch.distributed as dist
@@ -63,6 +64,7 @@ def run(rank, n_gpus, hps):
         logger = utils.get_logger(hps.model_dir)
         logger.info(hps)
         utils.check_git_hash(hps.model_dir)
+        wandb.init(project="Bert-VITS2", config=hps, sync_tensorboard=True)
         writer = SummaryWriter(log_dir=hps.model_dir)
         writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))