diff --git a/src/scheduler/Scheduler.cpp b/src/scheduler/Scheduler.cpp
index 242a2d0e6936bce1459946ee72e7d790d15c4fda..d456c8af2cfbfb9680069cda35eacc9941e9fe45 100644
--- a/src/scheduler/Scheduler.cpp
+++ b/src/scheduler/Scheduler.cpp
@@ -722,12 +722,12 @@ Aidge::MemoryManager Aidge::Scheduler::generateMemoryAutoConcat(bool incProducer
                             const auto parentOp = std::static_pointer_cast<OperatorTensor>(concatParent->getOperator());
                             const auto parentRequiredSize = parentOp->getRequiredMemory(outputIdx, {});
                             const auto parentOutputDims = (parentOp->getOutput(outputIdx)) ? parentOp->getOutput(outputIdx)->dims() : std::vector<DimSize_t>();
+                            const auto parentOutputFormat = (parentOp->getOutput(outputIdx)) ? parentOp->getOutput(outputIdx)->dataFormat() : DataFormat::Default;
 
                             // By default, specifies a fully monolithic memory block
                             std::size_t parentSize = parentRequiredSize.data;
 
-                            if (parentOutputDims.size() > 3) {
-                                // If it is possible, assume a NCHW layout
+                            if (parentOutputFormat == DataFormat::NHWC) {
                                 parentSize = parentOutputDims.end()[-3];
                             }