am5.scala

def endPoint(y: Variable): Receive = {
    case Forward(yHat) =>
      val l = loss(yHat, y)
      l.backward()
      wire.prev ! Backward(yHat.grad)
      trainingLoss += l.data.squeeze()

      if (trainingDataIterator.hasNext) {
        context become startPoint(TrainingDataIterator.next())
      } else {
        endOfEpoch()
      }
}