glenn-jocher commited on
Commit
afe1df3
1 Parent(s): 72d0614

dist.destroy_process_group() bug fix

Browse files
Files changed (1) hide show
  1. train.py +1 -1
train.py CHANGED
@@ -356,7 +356,7 @@ def train(hyp):
356
  if not opt.evolve:
357
  plot_results() # save as results.png
358
  print('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
359
- dist.destroy_process_group() if torch.cuda.device_count() > 1 else None
360
  torch.cuda.empty_cache()
361
  return results
362
 
 
356
  if not opt.evolve:
357
  plot_results() # save as results.png
358
  print('%g epochs completed in %.3f hours.\n' % (epoch - start_epoch + 1, (time.time() - t0) / 3600))
359
+ dist.destroy_process_group() if device.type != 'cpu' and torch.cuda.device_count() > 1 else None
360
  torch.cuda.empty_cache()
361
  return results
362