def get_sweep(hyper):
  """Sweeps over hyperparameters."""
  # Try sweep settings based on CIFAR finetuning sweep or older sweeps on
  # ViT-B/16 for EyePACS.
  step_sweep = hyper.zipit([
      # hyper.sweep('config.total_steps',
      #             [int(40_000 * x) for x in [0.5, 1.0, 1.5, 2.0]]),
      # hyper.sweep('config.lr.warmup_steps',
      #             [int(2000 * x) for x in [0.5, 1.0, 1.5, 2.0]]),
      hyper.sweep('config.total_steps',
                  [10_000, 15_000, 20_000, 25_000]),
      hyper.sweep('config.lr.warmup_steps',
                  [500, 750, 1000, 1250]),
  ])
  return hyper.product([
      # hyper.sweep('config.lr.base', [0.03/4, 0.01/4, 0.003/4, 0.001/4]),
      hyper.sweep('config.lr.base', [0.05, 0.03, 0.01, 0.005]),
      step_sweep,
      hyper.sweep('config.lr.decay_type', ['cosine', 'linear']),
      hyper.sweep('config.grad_clip_norm', [1.0, 2.5]),
      # Sweep on just APTOS as Severity has similar in-dist results.
      hyper.sweep('config.distribution_shift', ['aptos']),
      # hyper.sweep('config.distribution_shift', ['aptos', 'severity']),
  ])