From 23d574322f6e9dfc9477efb7dfafcc7d865aeb4c Mon Sep 17 00:00:00 2001
From: Giuseppe Franco <giuseppefranco4@gmail.com>
Date: Tue, 12 Sep 2023 16:56:15 +0100
Subject: [PATCH] Remove comments

---
 src/brevitas/graph/equalize.py  | 6 ++----
 src/brevitas/graph/test_unet.py | 1 -
 2 files changed, 2 insertions(+), 5 deletions(-)

diff --git a/src/brevitas/graph/equalize.py b/src/brevitas/graph/equalize.py
index 840cc6246..ce58a6d91 100644
--- a/src/brevitas/graph/equalize.py
+++ b/src/brevitas/graph/equalize.py
@@ -681,7 +681,6 @@ def cat_handler(graph_model: GraphModule, starting_node: Node, state: WalkRegion
     state.srcs.clear()
     state.sinks.clear()
     state.history.clear()
-    # state.srcs.add((starting_node.target, ()))
     state.srcs[starting_node.target] = _UNSUPPORTED_OP
     state.update_offset = True
     state.offset = 0
@@ -837,7 +836,8 @@ def _extract_regions(
                 module = get_module(graph_model, node.target)
                 weight = get_weight_source([module])
                 eq_indexes = EqualizationIndexes(0, weight[0].shape[0], 0)
-                state.srcs[node.target] = eq_indexes
+                full_source_name = node.target + '$' + str(eq_indexes)
+                state.srcs[full_source_name] = eq_indexes
             find_sinks(graph_model, node, state)
             if state.sinks and _UNSUPPORTED_OP not in state.sinks.keys(
             ) and _UNSUPPORTED_OP not in state.srcs.keys():
@@ -851,10 +851,8 @@ def _extract_regions(
                 else:
                     region = Region(srcs=sorted_srcs, sinks=sorted_sinks)
 
-                # new_region = (sorted_tuple_srcs, sorted_tuple_sinks)
                 if region not in regions:
                     regions.append(region)
-                    # regions_name.add(region_names)
     return regions
 
 
diff --git a/src/brevitas/graph/test_unet.py b/src/brevitas/graph/test_unet.py
index e476f8f91..198a6c344 100644
--- a/src/brevitas/graph/test_unet.py
+++ b/src/brevitas/graph/test_unet.py
@@ -44,7 +44,6 @@ def main(args):
 
     # Perform equalization
     model, regions = EqualizeGraph(iterations=20, return_regions=True).apply(model)
-    print(len(regions))
     # New output after equalization
     out = model(x)
     # for region in regions: