code: def imagetransformer_b12l_8h_b256_uncond_dr03_tpu():
    
    hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
    update_hparams_for_tpu(hparams)
    hparams.batch_size = 2
    hparams.num_heads = 8   # heads are expensive on tpu
    hparams.num_decoder_layers = 12
    hparams.block_length = 256
    hparams.hidden_size = 512
    hparams.filter_size = 2048
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.layer_prepostprocess_dropout = 0.3
    return hparams
docs: TPU related 12 layer 8 heads model.