def METHOD_NAME(workload: spec.Workload,
                  current_param_container: spec.ParameterContainer,
                  current_params_types: spec.ParameterTypeTree,
                  model_state: spec.ModelAuxiliaryState,
                  hyperparameters: spec.Hyperparameters,
                  batch: Dict[str, spec.Tensor],
                  loss_type: spec.LossType,
                  optimizer_state: spec.OptimizerState,
                  eval_results: List[Tuple[int, float]],
                  global_step: int,
                  rng: spec.RandomState) -> spec.UpdateReturn:
  """Return (updated_optimizer_state, updated_params)."""
  del current_params_types
  del eval_results
  del model_state
  del loss_type
  optimizer = optimizer_state['optimizer']
  optimizer.zero_grad()
  current_model = current_param_container
  (logits, logits_padding), _ = workload.model_fn(
      current_model,
      batch,
      None,
      spec.ForwardPassMode.TRAIN,
      rng,
      update_batch_norm=True)
  loss_dict = workload.loss_fn(batch['targets'], (logits, logits_padding))
  summed_loss = loss_dict['summed']
  n_valid_examples = loss_dict['n_valid_examples']
  if USE_PYTORCH_DDP:
    # Use dist_nn.all_reduce to ensure correct loss and gradient scaling.
    summed_loss = dist_nn.all_reduce(summed_loss)
    n_valid_examples = dist_nn.all_reduce(n_valid_examples)
  loss = summed_loss / n_valid_examples
  loss.backward()
  for g in optimizer.param_groups:
    g['lr'] = get_learning_rate(global_step, hyperparameters)
  if hasattr(hyperparameters, 'grad_clip'):
    torch.nn.utils.clip_grad_norm_(
        current_model.parameters(), max_norm=hyperparameters.grad_clip)
  optimizer.step()
  return optimizer_state, current_param_container, None