code: def imagetransformer_sep_channels_8l_tpu():
    
    hparams = imagetransformer_sep_channels_8l()
    update_hparams_for_tpu(hparams)
    hparams.batch_size = 4
    hparams.num_heads = 4   # heads are expensive on tpu
    hparams.shared_embedding_and_softmax_weights = False
    return hparams
docs: Hparams for training imagetransformer on tpu.