def train(
        self,
        train_data_loader: DataLoader,
        val_data_loader: Optional[DataLoader] = None,
        calib_data_loader: Optional[DataLoader] = None,
        fit_config: FitConfig = FitConfig(),
        calib_config: CalibConfig = CalibConfig(),
        map_fit_config: Optional[FitConfig] = None,
    ) -> Dict[str, Status]:
        """
        Train the probabilistic model. This involves fitting the posterior distribution and calibrating the
        probabilistic model. Calibration is performed only if (1) `calib_data_loader` is passed and (2) the
        probabilistic model contains any calibrator.

        Parameters
        ----------
        train_data_loader : DataLoader
            A training data loader.
        val_data_loader : DataLoader
            A validation data loader. This is used to validate both posterior fitting and calibration.
        calib_data_loader : DataLoader
            A calibration data loader. If this is not passed, no calibration is performed.
        fit_config : FitConfig
            An object to configure the posterior distribution fitting.
        calib_config : CalibConfig
            An object to configure the calibration.
        map_fit_config : Optional[FitConfig] = None
            An object to configure a preliminary posterior distribution fitting via the Maximum-A-Posteriori (MAP)
            method.
            The fit methods of several supported posterior approximations, like the ones of
            :class:`~fortuna.prob_model.posterior.swag.swag_posterior.SWAGPosterior` and
            :class:`~fortuna.prob_model.posterior.laplace.laplace_posterior.LaplacePosterior`, start from a preliminary
            run of MAP, which can be configured via this object. If the method does not start from MAP, this argument is
            ignored.

        Returns
        -------
        Dict[str, Status]
            Status objects for both posterior fitting and calibration.

        """
        logging.info("Fit the posterior distribution...")
        fit_status = self.posterior.fit(
            train_data_loader=train_data_loader,
            val_data_loader=val_data_loader,
            fit_config=fit_config,
            map_fit_config=map_fit_config,
        )

        calib_status = None
        if calib_data_loader:
            calib_status = self.calibrate(
                calib_data_loader=calib_data_loader,
                val_data_loader=val_data_loader,
                calib_config=calib_config,
            )
            logging.info("Calibration completed.")
        return dict(fit_status=fit_status, calib_status=calib_status)

        Train the probabilistic model. This involves fitting the posterior distribution and calibrating the
        probabilistic model. Calibration is performed only if (1) `calib_data_loader` is passed and (2) the
        probabilistic model contains any calibrator.

        Parameters
        ----------
        train_data_loader : DataLoader
            A training data loader.
        val_data_loader : DataLoader
            A validation data loader. This is used to validate both posterior fitting and calibration.
        calib_data_loader : DataLoader
            A calibration data loader. If this is not passed, no calibration is performed.
        fit_config : FitConfig
            An object to configure the posterior distribution fitting.
        calib_config : CalibConfig
            An object to configure the calibration.
        map_fit_config : Optional[FitConfig] = None
            An object to configure a preliminary posterior distribution fitting via the Maximum-A-Posteriori (MAP)
            method.
            The fit methods of several supported posterior approximations, like the ones of
            :class:`~fortuna.prob_model.posterior.swag.swag_posterior.SWAGPosterior` and
            :class:`~fortuna.prob_model.posterior.laplace.laplace_posterior.LaplacePosterior`, start from a preliminary
            run of MAP, which can be configured via this object. If the method does not start from MAP, this argument is
            ignored.

        Returns
        -------
        Dict[str, Status]
            Status objects for both posterior fitting and calibration.

        
fortuna/prob_model/base.py