From ea9c6765caffd39525262a72576f57ddaedf1f4b Mon Sep 17 00:00:00 2001 From: James Betker Date: Thu, 2 Jul 2020 15:11:21 -0600 Subject: [PATCH] Move train imports into init_dist --- codes/train.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/codes/train.py b/codes/train.py index 7157ddad..b2b87afc 100644 --- a/codes/train.py +++ b/codes/train.py @@ -7,8 +7,6 @@ import shutil from tqdm import tqdm import torch -#import torch.distributed as dist -#import torch.multiprocessing as mp from data.data_sampler import DistIterSampler import options.options as option @@ -19,6 +17,10 @@ from time import time def init_dist(backend='nccl', **kwargs): + # These packages have globals that screw with Windows, so only import them if needed. + import torch.distributed as dist + import torch.multiprocessing as mp + """initialization for distributed training""" if mp.get_start_method(allow_none=True) != 'spawn': mp.set_start_method('spawn')