def imagetransformer_ae_imagenet():
  """For 64x64 ImageNet. ~56M trainable variables."""
  hparams = imagetransformer_ae_cifar()
  hparams.max_length = int(64 * 64 * 3)
  hparams.img_len = 64
  hparams.num_heads = 4  # Heads are expensive on TPUs.
  # Reduce architecture from 32x32 CIFAR-10 in order to fit in memory.
  hparams.num_decoder_layers = 8
  hparams.num_compress_steps = 2
  return hparams
For 64x64 ImageNet. ~56M trainable variables.