def __call__(self, x: jnp.ndarray, train: bool = True) -> jnp.ndarray:
        """
        Block forward pass.

        Parameters
        ----------
        x: jnp.ndarray
            Block inputs.
        train: bool
            Whether the call is performed during training.

        Returns
        -------
        jnp.ndarray
            Block outputs.
        """
        dropout = self.dropout(rate=self.dropout_rate, broadcast_dims=(1, 2))

        y = self.norm(name="bn1")(x)
        y = nn.relu(y)
        if self.dropout_rate > 0.0:
            y = dropout(y, deterministic=not train)
        y = self.conv(self.filters, (3, 3), self.strides, name="conv1")(y)
        y = self.norm(name="bn2")(y)
        y = nn.relu(y)
        if self.dropout_rate > 0.0:
            y = dropout(y, deterministic=not train)
        y = self.conv(self.filters, (3, 3), name="conv2")(y)

        # Apply an up projection in case of channel mismatch
        if (x.shape[-1] != self.filters) or self.strides != (1, 1):
            x = self.conv(self.filters, (3, 3), self.strides)(x)
        return x + y

        Block forward pass.

        Parameters
        ----------
        x: jnp.ndarray
            Block inputs.
        train: bool
            Whether the call is performed during training.

        Returns
        -------
        jnp.ndarray
            Block outputs.
        
fortuna/model/wideresnet.py