better handle cases when there is no previous node

fastmachinelearning · Dec 4, 2024 · 915d2e1 · 915d2e1
1 parent 76d06e7
commit 915d2e1
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/hls4ml/model/optimizer/passes/batchnorm_opt.py b/hls4ml/model/optimizer/passes/batchnorm_opt.py
@@ -166,7 +166,7 @@ class FuseConsecutiveBatchNormalization(OptimizerPass):
     """
 
     def match(self, node):
-        prev_node = node.get_input_node(node.inputs[0])
+        prev_node = node.get_input_node()
         basic_match = (
             isinstance(node, BatchNormalization)
             and isinstance(prev_node, BatchNormalization)
@@ -194,7 +194,7 @@ def match(self, node):
             return False
 
     def transform(self, model, node):
-        prev_node = node.get_input_node(node.inputs[0])
+        prev_node = node.get_input_node()
 
         prev_map = prev_node.get_output_use_map()
         if len(prev_map[prev_node.outputs[0]]) > 1:

diff --git a/hls4ml/model/optimizer/passes/bn_fuse.py b/hls4ml/model/optimizer/passes/bn_fuse.py
@@ -18,7 +18,7 @@ class FuseBatchNormalization(OptimizerPass):
     """
 
     def match(self, node):
-        prev_node = node.get_input_node(node.inputs[0])
+        prev_node = node.get_input_node()
         basic_match = (
             isinstance(node, BatchNormalization)
             and isinstance(prev_node, (Dense, Conv1D, Conv2D))