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

        Parameters
        ----------
        x: Array
            Inputs.

        jnp.ndarray
            Output of the hidden layers.
        """
        x = nn.Conv(
            features=6,
            kernel_size=(5, 5),
            strides=(1, 1),
            padding="valid",
            dtype=self.dtype,
        )(x)
        x = nn.relu(x)
        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = nn.Conv(
            features=16,
            kernel_size=(5, 5),
            strides=(1, 1),
            padding="valid",
            dtype=self.dtype,
        )(x)
        x = nn.relu(x)
        x = nn.max_pool(x, window_shape=(2, 2), strides=(2, 2))
        x = x.reshape((x.shape[0], -1))  # flatten
        x = nn.Dense(features=120, dtype=self.dtype)(x)
        x = nn.relu(x)
        x = nn.Dense(features=84, dtype=self.dtype)(x)
        return x

        Forward pass.

        Parameters
        ----------
        x: Array
            Inputs.

        jnp.ndarray
            Output of the hidden layers.
        
fortuna/model/lenet.py