async function train(images, vaeOpts, savePath, logDir) {
  const encoderModel = encoder(vaeOpts);
  const decoderModel = decoder(vaeOpts);
  const vaeModel = vae(encoderModel, decoderModel);

  let summaryWriter;
  if (logDir != null) {
    console.log(`Logging loss values to ${logDir}.`);
    console.log(
        `Use the following command to start the tensorboard backend server:`);
    console.log(`  tensorboard --logdir ${logDir}`);
    summaryWriter = tf.node.summaryFileWriter(logDir);
  }

  console.log('\n** Train Model **\n');

  // Because we use a custom loss function, we will use optimizer.minimize
  // instead of the more typical model.fit. We thus need to define an optimizer
  // and manage batching the data ourselves.

  // Create the optimizer
  const optimizer = tf.train.adam();

  // Group the data into batches.
  const batches = _.chunk(images, batchSize);

  // Run the train loop.
  let step = 0;
  for (let i = 0; i < epochs; i++) {
    console.log(`\nEpoch #${i + 1} of ${epochs}\n`);
    for (let j = 0; j < batches.length; j++) {
      const currentBatchSize = batches[j].length
      const batchedImages = batchImages(batches[j]);

      const reshaped =
          batchedImages.reshape([currentBatchSize, vaeOpts.originalDim]);

      // This is the model optimization step. We make a prediction
      // compute loss and return it so that optimizer.minimize can
      // adjust the weights of the model.
      optimizer.minimize(() => {
        const outputs = vaeModel.apply(reshaped);
        const loss = vaeLoss(reshaped, outputs, vaeOpts);
        process.stdout.write('.');
        if (j % 50 === 0) {
          console.log('\nLoss:', loss.dataSync()[0]);
        }
        if (summaryWriter != null) {
          summaryWriter.scalar('loss', loss, step++);
        }

        return loss;
      });
      tf.dispose([batchedImages, reshaped]);
    }
    console.log('');
    // Generate a preview after each epoch
    await generate(decoderModel, vaeOpts.latentDim);
  }

  console.log('done training');
  saveDecoder(savePath, decoderModel);
  }
Train the auto encoder

@param {number[][]} images Flattened images for VAE training.
@param {object} vaeOpts Options for the VAE model, including the following
  fields:
  - originaDim {number} Length of the input flattened image.
  - intermediateDim {number} Number of units of the intermediate (i.e.,
    hidden) dense layer.
  - latentDim {number} Dimensionality of the latent space (i.e,. z-space).
@param {string} savePath Path to which the decoder part of the VAE model
  will be saved after training.
@param {string?} logDir Optional path to log directory. If specified, the
  batch-by-batch loss values will be logged to the directory during training,
  so that the training process can be monitored using TensorBoard.