code: def imagetransformer_ae_imagenet():
    
    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
docs: For 64x64 ImageNet. ~56M trainable variables.