Data parallel fix for pytorch#1857 (pytorch#1880)

* Data parallel fix for pytorch#1857
searches recursively for variable in input

* parallel_apply.py lint

Author: csarofeen

torch/nn/parallel/parallel_apply.py

@@ -3,6 +3,23 @@
 from torch.autograd import Variable
 
 
+def get_a_var(obj):
+    if isinstance(obj, Variable):
+        return obj
+
+    if isinstance(obj, list) or isinstance(obj, tuple):
+        results = map(get_a_var, obj)
+        for result in results:
+            if isinstance(result, Variable):
+                return result
+    if isinstance(obj, dict):
+        results = map(get_a_var, obj.items())
+        for result in results:
+            if isinstance(result, Variable):
+                return result
+    return None
+
+
 def parallel_apply(modules, inputs, kwargs_tup=None):
     assert len(modules) == len(inputs)
     if kwargs_tup:
@@ -17,9 +34,7 @@ def parallel_apply(modules, inputs, kwargs_tup=None):
     results = {}
 
     def _worker(i, module, input, kwargs, results, lock):
-        var_input = input
-        while not isinstance(var_input, Variable):
-            var_input = var_input[0]
+        var_input = get_a_var(input)
         try:
             with torch.cuda.device_of(var_input):
                 output = module(*input, **kwargs)