diff --git a/src/operator/Slice.cpp b/src/operator/Slice.cpp index 405b11be7f3d9674efa4a601a9ee01c3ae07519d..97ec0a5171a8f13fee0a93557b6831443f10713a 100644 --- a/src/operator/Slice.cpp +++ b/src/operator/Slice.cpp @@ -63,15 +63,31 @@ void Aidge::Slice_OpImpl::forward() { prodInputDims[nbDims - i] = prodInputDims[nbDims - i + 1] * inputDims[nbDims - i]; } - std::size_t j = 0; - std::size_t i = 0; - for (; j < prodSlicedDims[0];) { - op.getOutput(0)->getImpl()->copy(op.getInput(0)->getImpl()->rawPtr(beginning + i), 1, j); + std::size_t i = beginning; + std::size_t size = 0; + std::size_t offset = 0; + for (std::size_t j = 0; j < prodSlicedDims[0];) { + ++size; ++i; ++j; + bool newChunk = false; for (std::size_t idx = nbDims - 1; idx > 0; --idx) { - i += j % prodSlicedDims[idx] == 0 ? substractedDims[idx] * prodInputDims[idx + 1] : 0; + if (j % prodSlicedDims[idx] == 0) { + i += substractedDims[idx] * prodInputDims[idx + 1]; + newChunk = true; + } } + + if (newChunk) { + op.getOutput(0)->getImpl()->copy(op.getInput(0)->getImpl()->rawPtr(beginning), size, offset); + beginning = i; + offset += size; + size = 0; + } + } + + if (size > 0) { + op.getOutput(0)->getImpl()->copy(op.getInput(0)->getImpl()->rawPtr(beginning), size, offset); } }