updates

ultralytics · Mar 2, 2020 · 84371f6 · 84371f6
1 parent 7823473
commit 84371f6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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