def __call__(self, x: jnp.ndarray):
        """
        Forward pass.

        Parameters
        ----------
        x: jnp.ndarray
            Outputs of the deep feature extractor sub-network.

        Returns
        -------
        jnp.ndarray
            Model outputs.
        """
        x = nn.relu(x)
        x = nn.Dense(features=self.output_dim, dtype=self.dtype)(x)
        return x

        Forward pass.

        Parameters
        ----------
        x: jnp.ndarray
            Outputs of the deep feature extractor sub-network.

        Returns
        -------
        jnp.ndarray
            Model outputs.
        
fortuna/model/lenet.py