From dd7b1b22a3eb336853a08791148d7a7584de6cd3 Mon Sep 17 00:00:00 2001
From: Olivier BICHLER <olivier.bichler@cea.fr>
Date: Tue, 5 Mar 2024 22:06:37 +0100
Subject: [PATCH] Fixed batchSize for unidimensional inputs
---
src/operator/FCImpl.cpp | 5 +++--
1 file changed, 3 insertions(+), 2 deletions(-)
diff --git a/src/operator/FCImpl.cpp b/src/operator/FCImpl.cpp
index bc4a7a7c..99524590 100644
--- a/src/operator/FCImpl.cpp
+++ b/src/operator/FCImpl.cpp
@@ -57,9 +57,10 @@ void Aidge::FCImpl_cpu::forward()
const auto& input2 = std::static_pointer_cast<Tensor>(mOp.getRawInput(2))->refCastFrom(input2Fallback, *std::static_pointer_cast<Tensor>(mOp.getRawOutput(0)));
// Call kernel
+ const auto batchSize = (input0.dims().size() > 1) ? input0.dims()[0] : 1;
kernelFunc(dynamic_cast<const FC_Op&>(mOp).getStaticAttributes(),
- input0.dims()[0],
- input0.size() / input0.dims()[0],
+ batchSize,
+ input0.size() / batchSize,
input0.getImpl()->rawPtr(), input1.getImpl()->rawPtr(), input2.getImpl()->rawPtr(),
getCPUPtr(mOp.getRawOutput(0)));
}
--
GitLab