def transformer_tpu_with_conv():
  """Cut down on the number of heads, and use convs instead."""
  hparams = transformer_tpu()
  hparams.num_heads = 4  # Heads are expensive on TPUs.
  hparams.ffn_layer = "conv_relu_conv"
  return hparams
Cut down on the number of heads, and use convs instead.
tensor2tensor/models/transformer.py