Skip to content

Commit 4020f04

Browse files
committed
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
1 parent 9c25b7c commit 4020f04

File tree

1 file changed

+4
-1
lines changed

1 file changed

+4
-1
lines changed

Diff for: morph_net/framework/op_regularizer_manager.py

+4-1
Original file line numberDiff line numberDiff line change
@@ -599,7 +599,10 @@ def _dfs_for_source_ops(self, output_boundary, input_boundary=None):
599599
input_boundary = set(input_boundary)
600600
else:
601601
input_boundary = set()
602-
to_visit = list(output_boundary)
602+
if isinstance(output_boundary, list):
603+
to_visit = output_boundary
604+
else:
605+
to_visit = list(output_boundary)
603606
visited = set()
604607
while to_visit:
605608
# Get next op and mark as visited.

0 commit comments

Comments
 (0)