class_docstr: ResNetV1.5.
class_docstr_tok: ["ResNetV1.5","."]
code: def __call__(self, x, train: bool = True):
      conv = partial(self.conv, use_bias=False, dtype=self.dtype)
      norm = partial(
          nn.BatchNorm,
          use_running_average=not train,
          momentum=0.9,
          epsilon=1e-5,
          dtype=self.dtype,
          axis_name="batch",
      )
  
      x = conv(
          self.num_filters,
          (7, 7),
          (2, 2),
          padding=[(3, 3), (3, 3)],
          name="conv_init",
      )(x)
      x = norm(name="bn_init")(x)
      x = nn.relu(x)
      x = nn.max_pool(x, (3, 3), strides=(2, 2), padding="SAME")
      for i, block_size in enumerate(self.stage_sizes):
          for j in range(block_size):
              strides = (2, 2) if i > 0 and j == 0 else (1, 1)
              x = self.block_cls(
                  self.num_filters * 2**i,
                  strides=strides,
                  conv=conv,
                  norm=norm,
                  act=self.act,
              )(x)
      x = jnp.mean(x, axis=(1, 2))
      x = nn.Dense(self.num_classes, dtype=self.dtype)(x)
      x = jnp.asarray(x, self.dtype)
      return x
  
file_path: flax/examples/imagenet/models.py
imports: from functools import partial
  from typing import Any, Callable, Sequence, Tuple
  from flax import linen
  import jax.numpy
label_desc: Use image node_id 1 for calling the ResNet obj's underlying member method code with example usage: obj.__call__(x, train) and returns: x