def __init__(self, out_ch, stride=1, use_projection=False,
               activation=jax.nn.relu, which_norm=hk.BatchNorm,
               which_conv=hk.Conv2D, use_se=False, se_ratio=0.25,
               name=None):
    super().__init__(name=name)
    self.out_ch = out_ch
    self.stride = stride
    self.use_projection = use_projection
    self.activation = activation
    self.which_norm = which_norm
    self.which_conv = which_conv
    self.use_se = use_se
    self.se_ratio = se_ratio

    self.width = self.out_ch // 4

    self.bn0 = which_norm(name='bn0')
    self.conv0 = which_conv(self.width, kernel_shape=1, with_bias=False,
                            padding='SAME', name='conv0')
    self.bn1 = which_norm(name='bn1')
    self.conv1 = which_conv(self.width, stride=self.stride,
                            kernel_shape=3, with_bias=False,
                            padding='SAME', name='conv1')
    self.bn2 = which_norm(name='bn2')
    self.conv2 = which_conv(self.out_ch, kernel_shape=1, with_bias=False,
                            padding='SAME', name='conv2')
    if self.use_projection:
      self.conv_shortcut = which_conv(self.out_ch, stride=stride,
                                      kernel_shape=1, with_bias=False,
                                      padding='SAME', name='conv_shortcut')
    if self.use_se:
      self.se = base.SqueezeExcite(self.out_ch, self.out_ch, self.se_ratio)