From e029f8491de48fd59c484b587167343b05efdfa0 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 16:49:59 +0800
Subject: [PATCH] Update train_ms.py

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

diff --git a/train_ms.py b/train_ms.py
index 3e0a187..718e3a0 100644
--- a/train_ms.py
+++ b/train_ms.py
@@ -7,7 +7,8 @@ import torch
 from torch import nn, optim
 from torch.nn import functional as F
 from torch.utils.data import DataLoader
-from torch.utils.tensorboard import SummaryWriter
+#from torch.utils.tensorboard import SummaryWriter
+from tensorboardX import SummaryWriter
 import torch.multiprocessing as mp
 import torch.distributed as dist
 from torch.nn.parallel import DistributedDataParallel as DDP