     Context.local.learningPhase = .training
     var trainingLossSum: Float = 0
     var trainingBatchCount = 0
+    let trainingBatchCountTotal = dataset.trainingPhrases.count
     for phrase in dataset.trainingPhrases {
       let sentence = phrase.numericalizedText
       let (loss, gradients) = valueWithGradient(at: model) { model -> Tensor<Float> in

        return Tensor(loss, on: device)
      }

      trainingLossSum += loss.scalarized()
      trainingBatchCount += 1
      optimizer.update(&model, along: gradients)
      LazyTensorBarrier()
      if hasNaN(gradients) {

      }
    }

    let bpc = getBpc(loss: validationLossSum, characterCount: validationCharacterCount)
    let validationLoss = validationLossSum / Float(validationBatchCount)

    print(
