From 883863e86c4caa28e08186b91c7aa59b907051d9 Mon Sep 17 00:00:00 2001
From: auphelia <jakobapk@web.de>
Date: Wed, 24 Jun 2020 16:17:31 +0100
Subject: [PATCH] [Transform] Add QuantAvgPool to infer datalayout trafo
---
src/finn/transformation/infer_data_layouts.py | 2 +-
1 file changed, 1 insertion(+), 1 deletion(-)
diff --git a/src/finn/transformation/infer_data_layouts.py b/src/finn/transformation/infer_data_layouts.py
index 9ac75578f..e7a6b8823 100644
--- a/src/finn/transformation/infer_data_layouts.py
+++ b/src/finn/transformation/infer_data_layouts.py
@@ -38,7 +38,7 @@ def _dims_to_layout(model, node, ndims):
return DataLayout.NC
else:
if node.domain == "finn":
- if node.op_type == "MultiThreshold":
+ if node.op_type == "MultiThreshold" or node.op_type == "QuantAvgPool2d":
mt_inst = registry.getCustomOp(node)
layout = mt_inst.get_nodeattr("data_layout")
if layout == "NHWC" and ndims == 4:
--
GitLab