def __init__(
        self,
        model: nn.Module,
        prior: Prior = IsotropicGaussianPrior(),
        posterior_approximator: PosteriorApproximator = SWAGPosteriorApproximator(),
        output_calibrator: Optional[nn.Module] = ClassificationTemperatureScaler(),
        model_editor: Optional[ModelEditor] = None,
        seed: int = 0,
    ):
        r"""
        A probabilistic classifier class.

        Parameters
        ----------
        model : nn.Module
            A model describing the deterministic relation between inputs and outputs. The outputs must correspond to
            the logits of a softmax probability vector. The output dimension must be the same as the number of classes.
            Let :math:`x` be input variables and :math:`w` the random model parameters. Then the model is described by
            a function :math:`f(w, x)`, where each component of :math:`f` corresponds to one of the classes.
        prior : Prior
            A prior distribution object. The default is an isotropic standard Gaussian. Let :math:`w` be the random
            model parameters. Then the prior is defined by a distribution :math:`p(w)`.
        posterior_approximator : PosteriorApproximator
            A posterior approximation method. The default method is SWAG.
        output_calibrator : Optional[nn.Module]
            An output calibrator object. The default is temperature scaling for classification, which rescales the
            logits with a scalar temperature parameter. Given outputs :math:`o` of the model manager, the output
            calibrator is described by a function :math:`g(\phi, o)`, where `phi` are deterministic
            calibration parameters.
        model_editor : ModelEditor
            A model_editor objects. It takes the forward pass and transforms the outputs.
        seed: int
            A random seed.

        Attributes
        ----------
        model : nn.Module
            See `model` in `Parameters`.
        model_manager : ClassificationModelManager
            This object orchestrates the model's forward pass.
        output_calibrator : nn.Module
            See `output_calibrator` in `Parameters`.
        prob_output_layer : ClassificationProbOutputLayer
            This object characterizes the distribution of target variable given the calibrated outputs. It is defined
            by :math:`p(y|o)=\text{Categorical}(y|p=softmax(o))`,
            where :math:`o` denotes the calibrated outputs and :math:`y` denotes a target variable.
        likelihood : ClassificationLikelihood
            The likelihood function. This is defined by
            :math:`p(y|w, \phi, x) = \text{Categorical}(y|p=\text{softmax}(g(\phi, f(w, x)))`.
        prior : Prior
            See `prior` in `Parameters`.
        joint : Joint
            This object describes the joint distribution of the target variables and the random parameters
            given the input variables and the calibration parameters, that is :math:`p(y, w|x, \phi)`.
        posterior_approximator : PosteriorApproximator
            See `posterior_approximator` in `Parameters`.
        posterior : Posterior
            This is the posterior approximation of the random parameters given the training data and the
            calibration parameters, that is :math:`p(w|\mathcal{D}, \phi)`, where :math:`\mathcal{D}` denotes the
            training data set and :math:`\phi` the calibration parameters.
        predictive : ClassificationPredictive
            This denotes the predictive distribution, that is :math:`p(y|\phi, x, \mathcal{D})`. Its statistics are
            approximated via a Monte Carlo approach by sampling from the posterior approximation.
        """
        self.model = model
        self.prior = prior
        self.output_calibrator = output_calibrator

        self.output_calib_manager = OutputCalibManager(
            output_calibrator=output_calibrator
        )
        self.prob_output_layer = self._get_prob_output_layer(model)

        model_manager_cls = getattr(
            ClassificationModelManagers, posterior_approximator.__str__()
        ).value
        self.model_manager = self._get_model_manager(
            model, model_editor, model_manager_cls, posterior_approximator
        )

        self.likelihood = ClassificationLikelihood(
            self.model_manager, self.prob_output_layer, self.output_calib_manager
        )
        self.joint = Joint(self.prior, self.likelihood)

        self.posterior = getattr(
            PosteriorApproximations, posterior_approximator.__str__()
        ).value(joint=self.joint, posterior_approximator=posterior_approximator)
        self.predictive = ClassificationPredictive(self.posterior)

        super().__init__(seed=seed)

        A probabilistic classifier class.

        Parameters
        ----------
        model : nn.Module
            A model describing the deterministic relation between inputs and outputs. The outputs must correspond to
            the logits of a softmax probability vector. The output dimension must be the same as the number of classes.
            Let :math:`x` be input variables and :math:`w` the random model parameters. Then the model is described by
            a function :math:`f(w, x)`, where each component of :math:`f` corresponds to one of the classes.
        prior : Prior
            A prior distribution object. The default is an isotropic standard Gaussian. Let :math:`w` be the random
            model parameters. Then the prior is defined by a distribution :math:`p(w)`.
        posterior_approximator : PosteriorApproximator
            A posterior approximation method. The default method is SWAG.
        output_calibrator : Optional[nn.Module]
            An output calibrator object. The default is temperature scaling for classification, which rescales the
            logits with a scalar temperature parameter. Given outputs :math:`o` of the model manager, the output
            calibrator is described by a function :math:`g(\phi, o)`, where `phi` are deterministic
            calibration parameters.
        model_editor : ModelEditor
            A model_editor objects. It takes the forward pass and transforms the outputs.
        seed: int
            A random seed.

        Attributes
        ----------
        model : nn.Module
            See `model` in `Parameters`.
        model_manager : ClassificationModelManager
            This object orchestrates the model's forward pass.
        output_calibrator : nn.Module
            See `output_calibrator` in `Parameters`.
        prob_output_layer : ClassificationProbOutputLayer
            This object characterizes the distribution of target variable given the calibrated outputs. It is defined
            by :math:`p(y|o)=\text{Categorical}(y|p=softmax(o))`,
            where :math:`o` denotes the calibrated outputs and :math:`y` denotes a target variable.
        likelihood : ClassificationLikelihood
            The likelihood function. This is defined by
            :math:`p(y|w, \phi, x) = \text{Categorical}(y|p=\text{softmax}(g(\phi, f(w, x)))`.
        prior : Prior
            See `prior` in `Parameters`.
        joint : Joint
            This object describes the joint distribution of the target variables and the random parameters
            given the input variables and the calibration parameters, that is :math:`p(y, w|x, \phi)`.
        posterior_approximator : PosteriorApproximator
            See `posterior_approximator` in `Parameters`.
        posterior : Posterior
            This is the posterior approximation of the random parameters given the training data and the
            calibration parameters, that is :math:`p(w|\mathcal{D}, \phi)`, where :math:`\mathcal{D}` denotes the
            training data set and :math:`\phi` the calibration parameters.
        predictive : ClassificationPredictive
            This denotes the predictive distribution, that is :math:`p(y|\phi, x, \mathcal{D})`. Its statistics are
            approximated via a Monte Carlo approach by sampling from the posterior approximation.
        
fortuna/prob_model/classification.py