diff --git a/src/operator/DepthToSpace.cpp b/src/operator/DepthToSpace.cpp
index 32db101fc2bc013f0a832d0f5455e5f69cdebc63..374e7a1c7ac4f596bc88097f37a87b77b51c0b67 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>();