def __call__(self, x, *, train: bool):
    """Apply the ResNet to the inputs `x`.

    Args:
      x: Inputs.
      train: Whether to use BatchNorm in training or inference mode.

    Returns:
      pool: The intermediate output of shape (N, 7, 7, 2048).
      out: The output head with shape (N, num_classes).
    """
    width = 64 * self.width_factor
    norm = functools.partial(
        nn.BatchNorm, use_running_average=not train, momentum=0.9)

    # Root block
    x = nn.Conv(
        features=width,
        kernel_size=(7, 7),
        strides=(2, 2),
        use_bias=False,
        name="init_conv")(
            x)
    x = norm(name="init_bn")(x)
    x = nn.max_pool(x, (3, 3), strides=(2, 2), padding="SAME")

    # Stages
    for i, stage_size in enumerate(self.stage_sizes):
      x = ResNetStage(
          stage_size,
          filters=width * 2**i,
          block_cls=self.block_cls,
          norm=norm,
          first_block_strides=(1, 1) if i == 0 else (2, 2),
          name=f"stage{i + 1}")(
              x)

    # Head
    pool = x
    out = jnp.mean(pool, axis=(1, 2))
    out = nn.Dense(
        self.num_classes, kernel_init=nn.initializers.zeros, name="head")(out)
    return pool, out
Apply the ResNet to the inputs `x`.

    Args:
      x: Inputs.
      train: Whether to use BatchNorm in training or inference mode.

    Returns:
      pool: The intermediate output of shape (N, 7, 7, 2048).
      out: The output head with shape (N, num_classes).
    
videogvt/train_lib/resnet_v1.py