def check_training_step(steps_per_epoch_from_dataset: int,
                        config: ml_collections.ConfigDict):
  """Checks whether the step in learning rate and config matches."""
  lr_configs = config.lr_configs
  steps_per_cycle = lr_configs.get('steps_per_cycle')
  total_steps = lr_configs.get('total_steps')
  assert not steps_per_cycle or not total_steps, (
      f'Do not set both steps_per_cycle {total_steps} and '
      f'total_steps{total_steps}.')

  total_steps_in_lr = total_steps or steps_per_cycle
  assert total_steps_in_lr is not None and total_steps_in_lr > 0, (
      'Must define total_steps or steps_per_cycle.')
  if config.get('lr_configs.steps_per_epoch'):
    given_steps_per_epoch = config.lr_configs.steps_per_epoch
    if steps_per_epoch_from_dataset != given_steps_per_epoch:
      logging.warning(('Mismatched step_per_epoch: given=%d vs actual=%d '
                       'batch_size=%d'), given_steps_per_epoch,
                      steps_per_epoch_from_dataset, config.batch_size)
Checks whether the step in learning rate and config matches.
videogvt/train_lib/train_utils.py