function decodeZ(inputTensor) {
  return tf.tidy(() => {
    const res = decoder.predict(inputTensor).mul(255).cast('int32');
    const reshaped = res.reshape(
        [inputTensor.shape[0], IMAGE_HEIGHT, IMAGE_WIDTH, IMAGE_CHANNELS]);
    return reshaped;
  });
}
Decode a (batch of) z vector into an image tensor. Z is the vector in latent
space that we want to generate an image for.

Returns an image tensor of the shape [batch, IMAGE_HEIGHT, IMAGE_WIDTH,
IMAGE_CHANNELS]

@param {Tensor2D} inputTensor of shape [batch, LATENT_DIMS]