Update train.py

2022-11-17 11:31:43 +08:00 · 2022-11-17 11:31:43 +08:00 · 58f23dc85a
parent d430234d81
commit 58f23dc85a
1 changed files with 3 additions and 0 deletions
--- a/train.py
+++ b/train.py
@ -40,6 +40,9 @@ def train(model, local_rank):
    if local_rank == 0:
        writer = SummaryWriter('train')
        writer_val = SummaryWriter('validate')
+    else:
+        writer = None
+        writer_val = None
    step = 0
    nr_eval = 0
    dataset = VimeoDataset('train')