diff --git a/include/otbTensorflowMultisourceModelLearningBase.hxx b/include/otbTensorflowMultisourceModelLearningBase.hxx
index 43a3da188f693556d6a226d733ab243be487493a..55f0de33d555efe5f1a45528cb91bfa53c67d86f 100644
--- a/include/otbTensorflowMultisourceModelLearningBase.hxx
+++ b/include/otbTensorflowMultisourceModelLearningBase.hxx
@@ -138,7 +138,7 @@ TensorflowMultisourceModelLearningBase<TInputImage>
     // Batch start and size
     const IndexValueType sampleStart = batch * m_BatchSize;
     IndexValueType batchSize = m_BatchSize;
-    if (rest != 0)
+    if (rest != 0 && batch == nBatches - 1)
     {
       batchSize = rest;
     }