def transformer_cifar10_memory_v0():
  """HParams for training image_cifar10_plain_gen_flat_rev with memory."""
  hparams = transformer_wikitext103_l4k_memory_v0()

  hparams.num_hidden_layers = 6

  hparams.max_length = 32 * 32 * 3
  hparams.split_targets_chunk_length = 64 * 3
  hparams.split_targets_max_chunks = int(
      hparams.max_length / hparams.split_targets_chunk_length)
  hparams.num_memory_items = 128 * 3

  # Since this is an image problem, batch size refers to examples (not tokens)
  target_images_per_batch = 4
  hparams.batch_size = int(target_images_per_batch * (
      hparams.max_length / hparams.split_targets_chunk_length))

  # The recurrent memory needs to know the actual batch size (in sequences)
  hparams.recurrent_memory_batch_size = hparams.batch_size

  hparams.max_relative_position = (
      hparams.num_memory_items + hparams.split_targets_chunk_length)

  return hparams
HParams for training image_cifar10_plain_gen_flat_rev with memory.