From 84371f68117cae975eabfa78cdf8a2aa1b78e4ba Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Sun, 1 Mar 2020 21:33:16 -0800 Subject: [PATCH] updates --- train.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/train.py b/train.py index 1a32e3c0..88cf3c42 100644 --- a/train.py +++ b/train.py @@ -158,7 +158,7 @@ def train(): # plt.savefig('LR.png', dpi=300) # Initialize distributed training - if device.type != 'cpu' and torch.cuda.device_count() > 1: + if device.type != 'cpu' and torch.cuda.device_count() > 1 and torch.distributed.is_available(): dist.init_process_group(backend='nccl', # 'distributed backend' init_method='tcp://127.0.0.1:9999', # distributed training init method world_size=1, # number of nodes for distributed training