Update train_ms.py

2023-08-23 16:49:59 +08:00
parent 432e3923f6
commit e029f8491d
1 changed files with 2 additions and 1 deletions
--- 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