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

        Parameters
        ----------
        x: jnp.ndarray
            Deep feature extractor representation.
        train: bool
            Whether the call is performed during training.

        Returns
        -------
        jnp.ndarray
            Output of the subnetwork.
        """
        x = nn.Dense(self.output_dim, dtype=self.dtype)(x)
        x = jnp.asarray(x, self.dtype)
        return x

        Output subnetwork forward pass.

        Parameters
        ----------
        x: jnp.ndarray
            Deep feature extractor representation.
        train: bool
            Whether the call is performed during training.

        Returns
        -------
        jnp.ndarray
            Output of the subnetwork.
        
fortuna/model/resnet.py