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

FIX: correct batch size for last iteration

parent d2195d91
......@@ -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;
}
......
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