def __init__(self, width, num_classes,
               variant='ResNet50',
               which_norm='BatchNorm', norm_kwargs=None,
               activation='relu', drop_rate=0.0,
               fc_init=jnp.zeros, conv_kwargs=None,
               preactivation=True, use_se=False, se_ratio=0.25,
               name='ResNet'):
    super().__init__(name=name)
    self.width = width
    self.num_classes = num_classes
    self.variant = variant
    self.depth_pattern = self.variant_dict[variant]['depth']
    self.activation = getattr(jax.nn, activation)
    self.drop_rate = drop_rate
    self.which_norm = getattr(hk, which_norm)
    if norm_kwargs is not None:
      self.which_norm = functools.partial(self.which_norm, **norm_kwargs)
    if conv_kwargs is not None:
      self.which_conv = functools.partial(hk.Conv2D, **conv_kwargs)
    else:
      self.which_conv = hk.Conv2D
    self.preactivation = preactivation

    # Stem
    self.initial_conv = self.which_conv(16 * self.width, kernel_shape=7,
                                        stride=2, padding='SAME',
                                        with_bias=False, name='initial_conv')
    if not self.preactivation:
      self.initial_bn = self.which_norm(name='initial_bn')
    which_block = ResBlockV2 if self.preactivation else ResBlockV1
    # Body
    self.blocks = []
    for multiplier, blocks_per_stage, stride in zip([64, 128, 256, 512],
                                                    self.depth_pattern,
                                                    [1, 2, 2, 2]):
      for block_index in range(blocks_per_stage):
        self.blocks += [which_block(multiplier * self.width,
                                    use_projection=block_index == 0,
                                    stride=stride if block_index == 0 else 1,
                                    activation=self.activation,
                                    which_norm=self.which_norm,
                                    which_conv=self.which_conv,
                                    use_se=use_se,
                                    se_ratio=se_ratio)]

    # Head
    self.final_bn = self.which_norm(name='final_bn')
    self.fc = hk.Linear(self.num_classes, w_init=fc_init, with_bias=True)