From 4020f04c3d733e8c24b6f0c762d136a608a702ce Mon Sep 17 00:00:00 2001 From: The MorphNet Team Date: Thu, 12 Mar 2020 14:32:08 -0700 Subject: [PATCH] Fixing a small issue in OpRegularizerManager. Internally, OpRegularizerManager was passing output_boundary as a list to _dfs_for_source_ops, which was packaging it again as a list. The fix now allows _dfs_for_source_ops to support a single output_boundary node or a list of nodes. PiperOrigin-RevId: 300625458 --- morph_net/framework/op_regularizer_manager.py | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/morph_net/framework/op_regularizer_manager.py b/morph_net/framework/op_regularizer_manager.py index eb713e5..afa87fe 100644 --- a/morph_net/framework/op_regularizer_manager.py +++ b/morph_net/framework/op_regularizer_manager.py @@ -599,7 +599,10 @@ def _dfs_for_source_ops(self, output_boundary, input_boundary=None): input_boundary = set(input_boundary) else: input_boundary = set() - to_visit = list(output_boundary) + if isinstance(output_boundary, list): + to_visit = output_boundary + else: + to_visit = list(output_boundary) visited = set() while to_visit: # Get next op and mark as visited.