From b992df125412c94626ac33c2e7a8ac14b2d1d8f4 Mon Sep 17 00:00:00 2001
From: Tobi-Alonso <tobi.alonso@gmail.com>
Date: Tue, 7 Jul 2020 18:50:48 +0100
Subject: [PATCH] [Transform] Fix for small graph and do not absorb neg mul
---
src/finn/transformation/insert_topk.py | 3 +++
1 file changed, 3 insertions(+)
diff --git a/src/finn/transformation/insert_topk.py b/src/finn/transformation/insert_topk.py
index 70012c1bc..3ef6ef1b1 100644
--- a/src/finn/transformation/insert_topk.py
+++ b/src/finn/transformation/insert_topk.py
@@ -52,6 +52,7 @@ class InsertTopK(Transformation):
if test:
init = model.get_initializer(node.input[1])
test = test and (init is not None) and all(x == 1 for x in init.shape)
+ test = test and init > 0
return test
def apply(self, model):
@@ -73,6 +74,8 @@ class InsertTopK(Transformation):
model.graph.node.remove(final_node)
graph_out_name = model.graph.output[0].name
final_node = model.find_producer(graph_out_name)
+ if final_node is None:
+ break
out_shape = model.get_tensor_shape(graph_out_name)
out_dtype = model.get_tensor_datatype(graph_out_name)
--
GitLab