def get_dataset_config():
  """Config for dataset and input pipeline."""
  config = ml_collections.ConfigDict()
  config.data_dir = 'data/spatial_relation_splits'
  config.grid_size = 6
  config.max_seq_len = 15
  config.max_target_seq_len = 25
  config.num_epochs = 40
  config.img_dim = 16
  config.eos_idx = 2
  config.train_split = 'train'
  config.eval_split = 'dev'
  config.test_splits = [
      'test', 'visual', 'relation', 'referent', 'relative_position_1',
      'relative_position_2'
  ]
  config.split = 'spatial_relation_splits'
  # If batches should be added to evaluate the entire dataset.
  config.eval_pad_last_batch = True
  config.test_pad_last_batch = True
  config.shuffle_buffer_size = 1000
  config.train_per_device_batch_size = 128
  config.eval_per_device_batch_size = 1000
  config.test_per_device_batch_size = 1000

  return config
Config for dataset and input pipeline.
language/gscan/xattn_model/configs/spatial_relation.py