commit 74c834e0dc7b0253487d2aa620a0a90dfe8b8fb8
author Kevin Tse <ktse@fb.com> Tue Dec 21 15:17:40 2021 -0800
committer Facebook GitHub Bot <facebook-github-bot@users.noreply.github.com> Tue Dec 21 15:21:04 2021 -0800
tree d240fa881bc03a50e504f98b19c707f4afcd7ee9
parent 23902fb8955d48b8d1dfdd54c40e72c2744cac8e

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

Summary:
Pull Request resolved: https://github.com/pytorch/pytorch/pull/70214

cc VitalyFedyunin ejguan NivekT

Test Plan: Imported from OSS

Reviewed By: ejguan

Differential Revision: D33255306

Pulled By: NivekT

fbshipit-source-id: de2fe6bf08328e481c714aaad390db771073469e

torch/utils/data/graph.py

diff --git a/torch/utils/data/graph.py b/torch/utils/data/graph.py
index cc73fa4..786aa56 100644
--- a/torch/utils/data/graph.py
+++ b/torch/utils/data/graph.py
@@ -40,17 +40,17 @@
             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