device_ids can be None again in data_parallel (pytorch#1187)

Author: soumith

test/test_nn.py

@@ -996,6 +996,9 @@ def test_data_parallel(self):
         self.assertEqual(out.get_device(), 0)
         self.assertEqual(out.data, expected_out)
 
+        # Check for None device_ids
+        out = dp.data_parallel(l, i)
+
     @unittest.skipIf(not TEST_MULTIGPU, "multi-GPU not supported")
     def test_data_parallel_nested_output(self):
         def fn(input):

torch/nn/parallel/data_parallel.py

@@ -74,7 +74,7 @@ def gather(self, outputs, output_device):
         return gather(outputs, output_device, dim=self.dim)
 
 
-def data_parallel(module, inputs, device_ids, output_device=None, dim=0, module_kwargs=None):
+def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, module_kwargs=None):
     """Evaluates module(input) in parallel across the GPUs given in device_ids.
 
     This is the functional version of the DataParallel module.
@@ -92,6 +92,9 @@ def data_parallel(module, inputs, device_ids, output_device=None, dim=0, module_
     if not isinstance(inputs, tuple):
         inputs = (inputs,)
 
+    if device_ids is None:
+        device_ids = list(range(torch.cuda.device_count()))
+
     if output_device is None:
         output_device = device_ids[0]