From 94b70000e4e95b209c89e9d3f4cb97c1ad3e630c Mon Sep 17 00:00:00 2001
From: NAUD Maxence <maxence.naud@cea.fr>
Date: Wed, 10 Jul 2024 09:11:02 +0000
Subject: [PATCH] move depthtospace channel size / blocksize assertion to
 forwarddims

---
 src/operator/DepthToSpace.cpp | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/src/operator/DepthToSpace.cpp b/src/operator/DepthToSpace.cpp
index 32db101fc..374e7a1c7 100644
--- a/src/operator/DepthToSpace.cpp
+++ b/src/operator/DepthToSpace.cpp
@@ -26,8 +26,6 @@ void Aidge::DepthToSpace_OpImpl::forward() {
 
     // Get input dimensions
     const auto& dims = op.getInput(0)->dims<4>();
-    // Assert that c is divisible by blocksize squared
-    assert(dims[1] % (op.blockSize() * op.blockSize()) == 0 && "Number of channels must be divisible by blocksize squared");
     // get final output dimension
     const std::array<DimSize_t, 4> final_dims = op.getOutput(0)->dims<4>();
 
@@ -63,6 +61,7 @@ const std::string Aidge::DepthToSpace_Op::Type = "DepthToSpace";
 bool Aidge::DepthToSpace_Op::forwardDims(bool /*allowDataDependency*/) {
     if (inputsAssociated()) {
         AIDGE_ASSERT(getInput(0)->nbDims() == 4, "{} Operator only accepts 4-D input Tensors.", DepthToSpace_Op::Type);
+        AIDGE_ASSERT(getInput(0)->dims()[1] % (blockSize() * blockSize()) == 0, "Number of channels must be divisible by blocksize squared");
 
         // Compute output dims
         const std::array<DimSize_t, 4>& inDims = getInput(0)->dims<4>();
-- 
GitLab