   for (size_t numEpoch = 0; numEpoch < epochs_; ++numEpoch) {
     std::cerr << "Starting epoch #" << numEpoch << std::endl;
     size_t startId = 0;
-    size_t endId = startId + maxBatchSize_;
 
-    while (endId < numExamples) {
-      PrepareBatch(startId, endId, maxBatchSize_, shuffle, xt, yt);
+    while (startId < numExamples) {
+      size_t batchSize = std::min(maxBatchSize_, numExamples - startId);
+      size_t endId = startId + batchSize;
+
+      PrepareBatch(startId, endId, batchSize, shuffle, xt, yt);
       *inX_ = xt;
       *inY_ = yt;
 

      UpdateModel();

      startId += maxBatchSize_;
      endId += maxBatchSize_;
    }
  }
}

	 << " " << startYDataId << "-" << endYDataId
	 << endl;
    */

    std::copy(xData_.begin() + startXDataId,
        xData_.begin() + endXDataId,
        x.begin() + startXId);
