def __call__(self, x: Array, train: bool = True) -> jnp.ndarray:
        """
        Deep feature extractor subnetwork forward pass.

        Parameters
        ----------
        x: Array
            Input data.
        train: bool
            Whether the call is performed during training.

        Returns
        -------
        jnp.ndarray
            Deep feature extractor representation.
        """
        if hasattr(self, "spectral_norm"):
            conv = self.spectral_norm(self.conv, train=train)
        else:
            conv = self.conv

        blocks_per_group = (self.depth - 4) // 6

        dropout = self.dropout(rate=self.dropout_rate, broadcast_dims=(1, 2))
        conv = partial(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,
        )

        x = conv(16, (3, 3), name="init_conv")(x)
        if self.dropout_rate > 0.0:
            x = dropout(x, deterministic=not train)
        x = WideResnetGroup(
            conv=conv,
            norm=norm,
            activation=self.activation,
            blocks_per_group=blocks_per_group,
            filters=16 * self.widen_factor,
            strides=(1, 1),
            dropout=self.dropout,
            dropout_rate=self.dropout_rate,
        )(x, train=train)
        x = WideResnetGroup(
            conv=conv,
            norm=norm,
            activation=self.activation,
            blocks_per_group=blocks_per_group,
            filters=32 * self.widen_factor,
            strides=(2, 2),
            dropout=self.dropout,
            dropout_rate=self.dropout_rate,
        )(x, train=train)
        x = WideResnetGroup(
            conv=conv,
            norm=norm,
            activation=self.activation,
            blocks_per_group=blocks_per_group,
            filters=64 * self.widen_factor,
            strides=(2, 2),
            dropout=self.dropout,
            dropout_rate=self.dropout_rate,
        )(x, train=train)
        x = norm()(x)
        x = nn.relu(x)
        x = nn.avg_pool(x, (8, 8))
        x = x.reshape((x.shape[0], -1))
        return x

        Deep feature extractor subnetwork forward pass.

        Parameters
        ----------
        x: Array
            Input data.
        train: bool
            Whether the call is performed during training.

        Returns
        -------
        jnp.ndarray
            Deep feature extractor representation.
        
fortuna/model/wideresnet.py