def METHOD_NAME(workload: spec.Workload,
                         model_params: spec.ParameterContainer,
                         model_state: spec.ModelAuxiliaryState,
                         hyperparameters: spec.Hyperparameters,
                         rng: spec.RandomState) -> spec.OptimizerState:
  del model_state
  del rng
  batch_size = get_batch_size('imagenet_vit')
  base_lr = hyperparameters.learning_rate * batch_size / 1024.
  optimizer_state = {
      'optimizer':
          torch.optim.Adam(
              model_params.parameters(),
              lr=base_lr,
              betas=(hyperparameters.beta1, hyperparameters.beta2),
              eps=hyperparameters.epsilon),
  }
  steps_per_epoch = workload.num_train_examples // batch_size
  scheduler1 = LinearLR(
      optimizer_state['optimizer'],
      start_factor=1e-10,
      end_factor=1.,
      total_iters=hyperparameters.warmup_epochs * steps_per_epoch)
  cosine_epochs = max(
      hyperparameters.num_epochs - hyperparameters.warmup_epochs, 1)
  scheduler2 = CosineAnnealingLR(
      optimizer_state['optimizer'], T_max=cosine_epochs * steps_per_epoch)
  optimizer_state['scheduler'] = SequentialLR(
      optimizer_state['optimizer'],
      schedulers=[scheduler1, scheduler2],
      milestones=[hyperparameters.warmup_epochs * steps_per_epoch])
  return optimizer_state