Better fix

zou3519 · zou3519 · commit 81973ead52e3 · 2018-02-05T14:23:23.000-07:00
diff --git a/torch/nn/parallel/scatter_gather.py b/torch/nn/parallel/scatter_gather.py
@@ -10,22 +10,22 @@ def scatter(inputs, target_gpus, dim=0):
     references to objects that are not variables. Does not
     support Tensors.
     """
-    def scatter_var(obj):
-        return Scatter.apply(target_gpus, None, dim, obj)
-
     def scatter_map(obj):
         if isinstance(obj, Variable):
-            return scatter_var(obj)
+            return Scatter.apply(target_gpus, None, dim, obj)
         assert not torch.is_tensor(obj), "Tensors not supported in scatter."
         if isinstance(obj, tuple) and len(obj) > 0:
-            return list(zip(*map(scatter_var, obj)))
+            return list(zip(*map(scatter_map, obj)))
         if isinstance(obj, list) and len(obj) > 0:
-            return list(map(list, zip(*map(scatter_var, obj))))
+            return list(map(list, zip(*map(scatter_map, obj))))
         if isinstance(obj, dict) and len(obj) > 0:
-            return list(map(type(obj), zip(*map(scatter_var, obj.items()))))
+            return list(map(type(obj), zip(*map(scatter_map, obj.items()))))
         return [obj for targets in target_gpus]
 
-    return scatter_map(inputs)
+    try:
+        return scatter_map(inputs)
+    finally:
+        scatter_map = None
 
 
 def scatter_kwargs(inputs, kwargs, target_gpus, dim=0):
@@ -46,17 +46,15 @@ def gather(outputs, target_device, dim=0):
     Gathers variables from different GPUs on a specified device
       (-1 means the CPU).
     """
-    def gather_vars(outputs):
-        return Gather.apply(target_device, dim, *outputs)
-
     def gather_map(outputs):
         out = outputs[0]
         if isinstance(out, Variable):
-            return gather_vars(outputs)
+            return Gather.apply(target_device, dim, *outputs)
         if out is None:
             return None
-        # Assuming outputs is a iterable of iterables and not
-        # an iterable of iterables of iterables (or more)
-        return type(out)(map(gather_vars, zip(*outputs)))
+        return type(out)(map(gather_map, zip(*outputs)))
 
-    return gather_map(outputs)
+    try:
+        return gather_map(outputs)
+    finally:
+        gather_map = None
