[DataPipe] adding a finally statement to ensure hook is reset (#70214)

NivekT · facebook-github-bot · commit 74c834e0dc7b · 2021-12-21T15:21:04.000-08:00
Summary: Pull Request resolved: #70214 cc VitalyFedyunin ejguan NivekT Test Plan: Imported from OSS Reviewed By: ejguan Differential Revision: D33255306 Pulled By: NivekT fbshipit-source-id: de2fe6bf08328e481c714aaad390db771073469e
diff --git a/torch/utils/data/graph.py b/torch/utils/data/graph.py
@@ -40,17 +40,17 @@ def reduce_hook(obj):
             captured_connections.append(obj)
             return stub_unpickler, ()
 
-    # TODO(VitalyFedyunin):  Better do it as `with` context for safety
-    IterDataPipe.set_reduce_ex_hook(reduce_hook)
-    if exclude_primitive:
-        IterDataPipe.set_getstate_hook(getstate_hook)
     try:
+        IterDataPipe.set_reduce_ex_hook(reduce_hook)
+        if exclude_primitive:
+            IterDataPipe.set_getstate_hook(getstate_hook)
         p.dump(scan_obj)
     except AttributeError:  # unpickable DataPipesGraph
         pass  # TODO(VitalyFedyunin): We need to tight this requirement after migrating from old DataLoader
-    IterDataPipe.set_reduce_ex_hook(None)
-    if exclude_primitive:
-        IterDataPipe.set_getstate_hook(None)
+    finally:
+        IterDataPipe.set_reduce_ex_hook(None)
+        if exclude_primitive:
+            IterDataPipe.set_getstate_hook(None)
     return captured_connections
 
 
