Hparams for training imagetransformer on tpu.

def imagetransformer_sep_channels_8l_tpu():
  """Hparams for training imagetransformer on 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