code: def img2img_transformer_base_tpu():
    
    hparams = img2img_transformer_base()
    update_hparams_for_tpu(hparams)
    hparams.batch_size = 2
    hparams.num_heads = 4   # heads are expensive on tpu
    hparams.num_decoder_layers = 8
    hparams.num_encoder_layers = 4
    hparams.shared_embedding_and_softmax_weights = False
    return hparams
docs: Hparams for training img2img_transformer on tpu.