code: def imagetransformer_b10l_dr03_moe_tpu():
    
    hparams = imagetransformer_b10l_4h_big_uncond_dr03_tpu()
    update_hparams_for_tpu(hparams)
    hparams.batch_size = 4
    hparams.num_heads = 4   # heads are expensive on tpu
    hparams.num_decoder_layers = 10
    hparams.layer_preprocess_sequence = "none"
    hparams.layer_postprocess_sequence = "dan"
    hparams.ffn_layer = "local_moe_tpu"
    return hparams
docs: Moe tpu params.