 
         optimizer.update(&autoencoder, along: 𝛁model)
     }
+
+    var testLossSum: Float = 0
+    var testBatchCount = 0
+    for batch in dataset.validation {
+        let sampleImages = batch.data
+        let testImages = autoencoder(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 = OldFashionMNIST(batchSize: batchSize, flattening: true)
// An autoencoder.
var autoencoder = Sequential {
    // The encoder.

}
let optimizer = RMSProp(for: autoencoder)


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

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