Let parallel_apply accept arbitrary inputs

Author: apaszke

torch/nn/parallel/data_parallel.py

@@ -67,7 +67,7 @@ def scatter(self, inputs, kwargs, device_ids):
         return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
 
     def parallel_apply(self, replicas, inputs, kwargs):
-        return parallel_apply(replicas, inputs, kwargs)
+        return parallel_apply(replicas, inputs, kwargs, self.device_ids)
 
     def gather(self, outputs, output_device):
         return gather(outputs, output_device, dim=self.dim)
@@ -101,5 +101,5 @@ def data_parallel(module, inputs, device_ids=None, output_device=None, dim=0, mo
     if len(device_ids) == 1:
         return module(*inputs[0], **module_kwargs[0])
     replicas = replicate(module, device_ids[:len(inputs)])
-    outputs = parallel_apply(replicas, inputs, module_kwargs)
+    outputs = parallel_apply(replicas, inputs, module_kwargs, device_ids)
     return gather(outputs, output_device, dim)

torch/nn/parallel/distributed.py

@@ -164,7 +164,7 @@ def scatter(self, inputs, kwargs, device_ids):
         return scatter_kwargs(inputs, kwargs, device_ids, dim=self.dim)
 
     def parallel_apply(self, replicas, inputs, kwargs):
-        return parallel_apply(replicas, inputs, kwargs)
+        return parallel_apply(replicas, inputs, kwargs, self.device_ids)
 
     def gather(self, outputs, output_device):
         return gather(outputs, output_device, dim=self.dim)

torch/nn/parallel/parallel_apply.py

@@ -20,40 +20,46 @@ def get_a_var(obj):
     return None
 
 
-def parallel_apply(modules, inputs, kwargs_tup=None):
+def parallel_apply(modules, inputs, kwargs_tup=None, devices=None):
     assert len(modules) == len(inputs)
-    if kwargs_tup:
+    if kwargs_tup is not None:
         assert len(modules) == len(kwargs_tup)
     else:
         kwargs_tup = ({},) * len(modules)
-    # Fast track
-    if len(modules) == 1:
-        return (modules[0](*inputs[0], **kwargs_tup[0]), )
+    if devices is not None:
+        assert len(modules) == len(devices)
+    else:
+        devices = [None] * len(modules)
 
     lock = threading.Lock()
     results = {}
 
-    def _worker(i, module, input, kwargs, results, lock):
-        var_input = get_a_var(input)
+    def _worker(i, module, input, kwargs, results, lock, device=None):
+        if device is None:
+            device = get_a_var(input).get_device()
         try:
-            with torch.cuda.device_of(var_input):
+            with torch.cuda.device(device):
                 output = module(*input, **kwargs)
             with lock:
                 results[i] = output
         except Exception as e:
             with lock:
                 results[i] = e
 
-    threads = [threading.Thread(target=_worker,
-                                args=(i, module, input, kwargs, results, lock),
-                                )
-               for i, (module, input, kwargs) in
-               enumerate(zip(modules, inputs, kwargs_tup))]
+    if len(modules) > 1:
+        threads = [threading.Thread(target=_worker,
+                                    args=(i, module, input, kwargs, results, lock, device),
+                                    )
+                   for i, (module, input, kwargs, device) in
+                   enumerate(zip(modules, inputs, kwargs_tup, devices))]
+
+        for thread in threads:
+            thread.start()
+        for thread in threads:
+            thread.join()
+    else:
+        _worker(0, modules[0], inputs[0], kwargs_tup[0], results, lock, devices[0])
 
-    for thread in threads:
-        thread.start()
-    for thread in threads:
-        thread.join()
     outputs = []
     for i in range(len(inputs)):
         output = results[i]