def train_step(
    train_state: TrainState,
    batch: Batch,
    *,
    model_dict: Dict[str, nn.Module],
    config: ml_collections.ConfigDict) -> Tuple[Any, metrics.Collection]:
  """Runs a single step of training.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where
      batch['inputs'].shape= device_bs, t, h, w, 3
    model_dict: A dictionary of generator and discriminator Flax models
    config: Configurations of the experiment.

  Returns:
    Updated state of training, computed metrics.
  """
  if config.vqgan.model_type == '2D' and config.get('dataset_type',
                                                    'video') == 'video':
    sampled_frames = video_utils.sample_frames_uniformly(
        batch['inputs'], n_sampled_frames=config.num_train_sampled_frames)
    batch = dict(inputs=sampled_frames)

  new_train_state, metrics_update = train_g_d(
      train_state,
      batch,
      model_dict=model_dict,
      config=config)
  return new_train_state, metrics_update
Runs a single step of training.

  Args:
    train_state: The state of training TrainState
    batch: A single batch of data. Dictionary where
      batch['inputs'].shape= device_bs, t, h, w, 3
    model_dict: A dictionary of generator and discriminator Flax models
    config: Configurations of the experiment.

  Returns:
    Updated state of training, computed metrics.
  
videogvt/trainers/vqgan_trainer.py