def afx_small_bfloat16():
  """Small transformer model with small batch size for fast step times."""
  hparams = afx_small()
  hparams.weight_dtype = "bfloat16"
  hparams.activation_dtype = "bfloat16"
  return hparams
Small transformer model with small batch size for fast step times.
tensor2tensor/models/research/adafactor_experiments.py