Commit 252f6a72 authored by remi cresson's avatar remi cresson
Browse files

FIX: correct batch size for last iteration

parent d2195d91
...@@ -138,7 +138,7 @@ TensorflowMultisourceModelLearningBase<TInputImage> ...@@ -138,7 +138,7 @@ TensorflowMultisourceModelLearningBase<TInputImage>
// Batch start and size // Batch start and size
const IndexValueType sampleStart = batch * m_BatchSize; const IndexValueType sampleStart = batch * m_BatchSize;
IndexValueType batchSize = m_BatchSize; IndexValueType batchSize = m_BatchSize;
if (rest != 0) if (rest != 0 && batch == nBatches - 1)
{ {
batchSize = rest; batchSize = rest;
} }
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment