Set of hyperparameters for a very small imagetransformer with MoE.

def imagetransformer_moe_tiny():
  """Set of hyperparameters for a very small imagetransformer with MoE."""
  hparams = imagetransformer_tiny()
  hparams.hidden_size = 64
  hparams.batch_size = 1
  hparams.num_hidden_layers = 3
  hparams.dec_attention_type = cia.AttentionType.MOE_LOCAL_1D
  hparams.add_hparam("moe_layers_decoder", "1")  # Which layer is MoE.
  hparams.moe_hidden_sizes = "1024"  # Hidden layer sizes (comma-separated).
  hparams.moe_num_experts = 16  # Number of experts in each MoE layer.
  hparams.moe_k = 2  # How many experts to use per batch element (try 2 or 4).
  hparams.moe_loss_coef = 1e-2  # MoE loss coefficient (1e-2 is usually ok).
  return hparams