enabling gpu size = 1 to run without data parallel

This commit is contained in:
William Falcon
2019-07-13 22:16:10 -04:00
parent 7da82c2560
commit e9f5913dac
+2 -2
View File
@@ -96,7 +96,7 @@ class Trainer(TrainerIO):
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = ','.join([str(x) for x in self.data_parallel_device_ids])
self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 0
self.data_parallel = self.data_parallel_device_ids is not None and len(self.data_parallel_device_ids) > 1
# process info
self.proc_rank = 0
@@ -295,7 +295,7 @@ class Trainer(TrainerIO):
# when using gpus, first thing we do is spawn a new process between each worker
# applies to single gpu, multi-gpu and multi-nodes
if self.on_gpu and len(self.data_parallel_device_ids) > 1:
if self.data_parallel:
self.experiment = self.experiment.get_meta_copy()
mp.spawn(self.dp_train, nprocs=len(self.data_parallel_device_ids), args=(model, ))
else: