From fd5754e4afc4050ea85bc084d92d6c490c03d229 Mon Sep 17 00:00:00 2001
From: NAUD Maxence <maxence.naud@cea.fr>
Date: Tue, 5 Nov 2024 15:24:45 +0000
Subject: [PATCH] [fix] Remove automatic squeeze in GlobalAvgPool forwarddims
---
src/operator/GlobalAveragePooling.cpp | 6 ++++--
1 file changed, 4 insertions(+), 2 deletions(-)
diff --git a/src/operator/GlobalAveragePooling.cpp b/src/operator/GlobalAveragePooling.cpp
index bbcfd0d28..57886ec2f 100644
--- a/src/operator/GlobalAveragePooling.cpp
+++ b/src/operator/GlobalAveragePooling.cpp
@@ -44,8 +44,10 @@ bool Aidge::GlobalAveragePooling_Op::forwardDims(bool /*allowDataDependency*/) {
// Global average pooling takes each filter, averages its values and uses
// it as an output(Much like a fancier flatten). 1st dim is batch 2nd is
// number of filter
- mOutputs[0]->resize({getInput(0)->dims().at(0),
- getInput(0)->dims().at(1)});
+ std::vector<DimSize_t> outputDims(getInput(0)->nbDims(), 1);
+ outputDims[0] = getInput(0)->dims()[0];
+ outputDims[1] = getInput(0)->dims()[1];
+ mOutputs[0]->resize(outputDims);
return true;
}
--
GitLab