From 7a694e7c2b821afdff59a34c05df559273caf248 Mon Sep 17 00:00:00 2001 From: Kevin Tse Date: Mon, 20 Dec 2021 18:53:47 -0500 Subject: [PATCH] [DataPipe] adding a finally statement to ensure hook is reset [ghstack-poisoned] --- torch/utils/data/graph.py | 14 +++++++------- 1 file changed, 7 insertions(+), 7 deletions(-) diff --git a/torch/utils/data/graph.py b/torch/utils/data/graph.py index cc73fa47bb94..786aa563e7c2 100644 --- 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