Skip to content

Commit 38ae1ae

Browse files
jmduartethesps
authored andcommitted
fix for linear activation -> split to multiple layers
1 parent 03fd5d4 commit 38ae1ae

File tree

1 file changed

+7
-6
lines changed

1 file changed

+7
-6
lines changed

hls4ml/model/graph.py

Lines changed: 7 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -454,13 +454,14 @@ def remove_node(self, node, rewire=True):
454454
if len(node.inputs) > 1 or len(node.outputs) > 1:
455455
raise Exception('Cannot rewire a node with multiple inputs/outputs')
456456
prev_node = self.graph.get(node.inputs[0])
457-
next_node = next((x for x in self.graph.values() if node.outputs[0] in x.inputs), None)
457+
next_nodes = [x for x in self.graph.values() if node.outputs[0] in x.inputs]
458458
if prev_node is not None:
459-
if next_node is not None:
460-
for i,_ in enumerate(next_node.inputs):
461-
if node.outputs[0] == next_node.inputs[i]:
462-
next_node.inputs[i] = prev_node.outputs[0]
463-
break
459+
if len(next_nodes) > 0:
460+
for next_node in next_nodes:
461+
for i,_ in enumerate(next_node.inputs):
462+
if node.outputs[0] == next_node.inputs[i]:
463+
next_node.inputs[i] = prev_node.outputs[0]
464+
break
464465
else:
465466
if not node.outputs[0] in self.outputs:
466467
raise Exception('Cannot rewire a node without child')

0 commit comments

Comments
 (0)