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.