   }
 
   guard let validationDataset = dataset.validation else {
+    let trainingLoss = trainingLossSum / Float(trainingBatchCount)
+
     print(
       """
       [Epoch \(epoch)] \
-      Training loss: \(trainingLossSum / Float(trainingBatchCount))
+      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
   }
 

  }

  let bpc = validationLossSum / Float(validationCharacterCount) / log(2)

  print(
    """
    [Epoch \(epoch)] \
    Bits per character: \(bpc) \
    Validation loss: \(validationLossSum / Float(validationBatchCount))
    """
  )
}

func hasNaN<T: KeyPathIterable>(_ t: T) -> Bool {

// Training flags
let dropoutProb = 0.5  // Dropout rate.
let order = 5  // Power of length penalty.
let maxEpochs = 1000  // Maximum number of training epochs.
var lossHistory = Array<Float>() // Keep track of loss.
var noImprovements = 0 // Consecutive epochs without improvements to loss.
let lambd: Float = 0.00075  // Weight of length penalty.
// Lexicon flags.
let maxLength = 10  // Maximum length of a string.
