TPU config for cifar 10.

def imagetransformer_b12l_4h_b128_uncond_dr03_tpu():
  """TPU config for cifar 10."""
  hparams = imagetransformer_bas8l_8h_big_uncond_dr03_imgnet()
  update_hparams_for_tpu(hparams)
  hparams.batch_size = 2
  hparams.num_heads = 4   # heads are expensive on tpu
  hparams.num_decoder_layers = 12
  hparams.block_length = 128
  hparams.hidden_size = 256
  hparams.filter_size = 2048
  hparams.layer_preprocess_sequence = "none"
  hparams.layer_postprocess_sequence = "dan"
  hparams.layer_prepostprocess_dropout = 0.1
  hparams.optimizer = "Adafactor"
  hparams.learning_rate_schedule = "rsqrt_decay"
  hparams.learning_rate_warmup_steps = 10000
  return hparams