fix corner case in kwargs for DataParallel (pytorch#930)

Author: soumith

test/test_nn.py

@@ -869,6 +869,9 @@ def local_test(out):
         out = dp.data_parallel(m, (var1, var2, float1), (0, 1))
         local_test(out)
 
+        out = dp.data_parallel(m, (var1, var2, float1), (1, 0))
+        local_test(out)
+
         out = dp.data_parallel(m, (var1, var2, float1), (0,))
         local_test(out)
 

torch/nn/parallel/data_parallel.py

@@ -78,17 +78,16 @@ def _to_cuda(obj):
 
         replicas = self.replicate(self.module, self.device_ids)
         scattered = self.scatter(inputs, self.device_ids)
-
+        used_gpus = len(scattered)  # The last GPU might not be used. For example, input of size 4, on 5 GPUs
         gpu_dicts = None
-        if kwargs:
-            scatter_kwargs = {}
+        if kwargs is not None:
+            gpu_dicts = [{} for i in range(used_gpus)]
             for key in kwargs.keys():
-                scatter_kwargs[key] = self.scatter(
-                    _to_cuda(kwargs[key]), self.device_ids)
-            gpu_dicts = tuple(
-                {key: values[i] for key, values in scatter_kwargs.items()}
-                for i in self.device_ids
-            )
+                scattered_kwargs = self.scatter(_to_cuda(kwargs[key]), self.device_ids)
+                assert len(scattered_kwargs) == used_gpus
+                for i in range(used_gpus):
+                    gpu_dicts[i][key] = scattered_kwargs[i]
+
         replicas = replicas[:len(scattered)]
         outputs = self.parallel_apply(replicas, scattered, gpu_dicts)
         return self.gather(outputs, self.output_device)