 
         optimizer.update(&model, along: 𝛁model)
     }
+
+    Context.local.learningPhase = .inference
+    var testLossSum: Float = 0
+    var testBatchCount = 0
+    for batch in dataset.validation {
+        let sampleImages = batch.data
+        let testImages = model(sampleImages)
+
+        do {
+            try saveImage(
+                sampleImages[0..<1], shape: (imageWidth, imageHeight), format: .grayscale,
+                directory: outputFolder, name: "epoch-\(epoch)-input")
+            try saveImage(
+                testImages[0..<1], shape: (imageWidth, imageHeight), format: .grayscale,
+                directory: outputFolder, name: "epoch-\(epoch)-output")
+        } catch {
+            print("Could not save image with error: \(error)")
+        }
+
+        testLossSum += meanSquaredError(predicted: testImages, expected: sampleImages).scalarized()
+        testBatchCount += 1
+    }
+
+    print(
+        """
+        [Epoch \(epoch)] \
+        Loss: \(testLossSum / Float(testBatchCount))
+        """
+    )
 }

let imageWidth = 28

let outputFolder = "./output/"
let dataset = OldKuzushijiMNIST(batchSize: batchSize, flattening: true)

// An autoencoder.
struct Autoencoder2D: Layer {


    @differentiable
    func callAsFunction(_ input: Tensor<Float>) -> Tensor<Float> {
        let resize = input.reshaped(to: [batchSize, 28, 28, 1])
        let encoder = resize.sequenced(through: encoder1,
            encoder2, encoder3, encoder4, encoder5, encoder6)
        let decoder = encoder.sequenced(through: decoder1,
            decoder2, decoder3, decoder4, decoder5, decoder6)
        return output(decoder).reshaped(to: [batchSize, imageHeight * imageWidth])
    }
}

var model = Autoencoder2D()
let optimizer = AdaDelta(for: model)

// Training loop
for (epoch, epochBatches) in dataset.training.prefix(epochCount).enumerated() {
    Context.local.learningPhase = .training
    for batch in epochBatches {
        let x = batch.data

        let 𝛁model = TensorFlow.gradient(at: model) { model -> Tensor<Float> in
            let image = model(x)
