   }
 
   let bpc = validationLossSum / Float(validationCharacterCount) / log(2)
+  let validationLoss = validationLossSum / Float(validationBatchCount)
 
   print(
     """
     [Epoch \(epoch)] \
     Bits per character: \(bpc) \
-    Validation loss: \(validationLossSum / Float(validationBatchCount))
+    Validation loss: \(validationLoss)
     """
   )
+
+  // Stop training when loss stops improving.
+  lossHistory.append(validationLoss)
+  if lossHistory.min() == validationLoss {
+    noImprovements = 0
+  } else {
+    noImprovements += 1
+    if noImprovements >= 5 { break }
+  }
 }
 
 func hasNaN<T: KeyPathIterable>(_ t: T) -> Bool {

// Training flags
let dropoutProb = 0.5  // Dropout rate.
let order = 5  // Power of length penalty.
let maxEpochs = 10  // Maximum number of training epochs.
let lambd: Float = 0.00075  // Weight of length penalty.
// Lexicon flags.
let maxLength = 10  // Maximum length of a string.

  }

  guard let validationDataset = dataset.validation else {
    let trainingLoss = trainingLossSum / Float(trainingBatchCount)

    print(
      """
      [Epoch \(epoch)] \
      Training loss: \(trainingLoss))
      """
    )

    // Stop training when loss stops improving.
    lossHistory.append(trainingLoss)
    if lossHistory.min() == trainingLoss {
      noImprovements = 0
    } else {
      noImprovements += 1
      if noImprovements >= 5 { break }
    }

    continue
  }

