def map_keys(batch):
  """Dataset returns 'image' and 'targets'. We want 'inputs' and 'label'."""
  batch['inputs'] = batch['image']
  batch.pop('image')
  batch['label'] = batch['targets']
  return batch