From 196e6cd246523ec7f75aed03ec18f55713d459fb Mon Sep 17 00:00:00 2001
From: Marcus Edel <marcus.edel@fu-berlin.de>
Date: Mon, 11 Jun 2018 15:57:27 +0200
Subject: [PATCH] Only increase the batch size if there are more samples left.

---
 src/mlpack/core/optimizers/bigbatch_sgd/bigbatch_sgd_impl.hpp | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/mlpack/core/optimizers/bigbatch_sgd/bigbatch_sgd_impl.hpp b/src/mlpack/core/optimizers/bigbatch_sgd/bigbatch_sgd_impl.hpp
index ca0855a458..4375e580c4 100644
--- a/src/mlpack/core/optimizers/bigbatch_sgd/bigbatch_sgd_impl.hpp
+++ b/src/mlpack/core/optimizers/bigbatch_sgd/bigbatch_sgd_impl.hpp
@@ -143,7 +143,7 @@ double BigBatchSGD<UpdatePolicyType>::Optimize(
         if (batchOffset <= 0)
           batchOffset = 1;
 
-        if (batchSize + batchOffset >= numFunctions)
+        if ((currentFunction + batchSize + batchOffset) >= numFunctions)
           break;
 
         // Update the stochastic gradient estimation.