 let imageWidth = 28
 
 let outputFolder = "./output/"
-let dataset = OldKuzushijiMNIST(batchSize: batchSize, flattening: true)
+let dataset = KuzushijiMNIST(batchSize: batchSize, device: Device.default, 
+    entropy: SystemRandomNumberGenerator(), flattening: true)
 
 // An autoencoder.
 struct Autoencoder2D: Layer {


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


    @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)
