Update train_ms.py
This commit is contained in:
@@ -8,6 +8,7 @@ from torch import nn, optim
|
|||||||
from torch.nn import functional as F
|
from torch.nn import functional as F
|
||||||
from torch.utils.data import DataLoader
|
from torch.utils.data import DataLoader
|
||||||
#from torch.utils.tensorboard import SummaryWriter
|
#from torch.utils.tensorboard import SummaryWriter
|
||||||
|
import wandb
|
||||||
from wandb.tensorboard import WandbSummaryWriter as SummaryWriter
|
from wandb.tensorboard import WandbSummaryWriter as SummaryWriter
|
||||||
import torch.multiprocessing as mp
|
import torch.multiprocessing as mp
|
||||||
import torch.distributed as dist
|
import torch.distributed as dist
|
||||||
@@ -63,6 +64,7 @@ def run(rank, n_gpus, hps):
|
|||||||
logger = utils.get_logger(hps.model_dir)
|
logger = utils.get_logger(hps.model_dir)
|
||||||
logger.info(hps)
|
logger.info(hps)
|
||||||
utils.check_git_hash(hps.model_dir)
|
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 = SummaryWriter(log_dir=hps.model_dir)
|
||||||
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
|
writer_eval = SummaryWriter(log_dir=os.path.join(hps.model_dir, "eval"))
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user