def create_train_state(config: ml_collections.ConfigDict, rng: np.ndarray,
                       input_shape: Sequence[int],
                       num_classes: int) -> Tuple[nn.Module, ModelState]:
  """Create and initialize the model.

  Args:
    config: Configuration for model.
    rng: JAX PRNG Key.
    input_shape: Shape of the inputs fed into the model.
    num_classes: Number of classes in the output layer.

  Returns:
    model: Flax nn.Nodule model architecture.
    state: The initialized ModelState with the optimizer.
  """
  if config.model_name == "resnet50":
    model_cls = resnet_v1.ResNet50
  else:
    raise ValueError(f"Model {config.model_name} not supported.")
  model = model_cls(num_classes=num_classes)
  variables = model.init(rng, jnp.ones(input_shape), train=False)
  params = variables["params"]
  batch_stats = variables["batch_stats"]
  return model, ModelState(params=params, batch_stats=batch_stats)
Create and initialize the model.

  Args:
    config: Configuration for model.
    rng: JAX PRNG Key.
    input_shape: Shape of the inputs fed into the model.
    num_classes: Number of classes in the output layer.

  Returns:
    model: Flax nn.Nodule model architecture.
    state: The initialized ModelState with the optimizer.
  
videogvt/train_lib/pretrained_model_utils.py