 }
 let optimizer = RMSProp(for: autoencoder)
 
-let individualTestImages = Batcher(on: dataset.test.dataset, batchSize: 1)
-var testImageIterator = individualTestImages.sequenced()
 
 // Training loop
-for epoch in 1...epochCount {
-    if let nextIndividualImage = testImageIterator.next() {
-        let sampleTensor = nextIndividualImage.first
-        let sampleImage = Tensor(
-            shape: [1, imageHeight * imageWidth], scalars: sampleTensor.scalars)
-
-        let testImage = autoencoder(sampleImage)
-
-        do {
-            try saveImage(
-                sampleImage, shape: (imageWidth, imageHeight), format: .grayscale,
-                directory: outputFolder, name: "epoch-\(epoch)-input")
-            try saveImage(
-                testImage, shape: (imageWidth, imageHeight), format: .grayscale,
-                directory: outputFolder, name: "epoch-\(epoch)-output")
-        } catch {
-            print("Could not save image with error: \(error)")
-        }
-
-        let sampleLoss = meanSquaredError(predicted: testImage, expected: sampleImage)
-        print("[Epoch: \(epoch)] Loss: \(sampleLoss)")
-    }
-
-    for batch in dataset.training.sequenced() {
-        let x = batch.first
+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)


        optimizer.update(&autoencoder, along: 𝛁model)
    }
}

let imageWidth = 28

let outputFolder = "./output/"
let dataset = FashionMNIST(batchSize: batchSize, device: Device.default, 
    entropy: SystemRandomNumberGenerator(), flattening: true)
// An autoencoder.
var autoencoder = Sequential {
    // The encoder.
