def compute_loss(
        self,
        model_output: torch.Tensor,
        target: torch.Tensor,
        sample_weight: torch.Tensor,
    ):
        """
        Computes a loss from a `model_output`, which represents the parameters of a given probability
        distribution for every ground truth value in `target`, and the `target` itself.
        """
        params_out = self._params_from_output(model_output)
        loss = self._nllloss(params_out, target, sample_weight)

        prior_params = self._prior_params
        use_prior = prior_params is not None and any(
            p is not None for p in prior_params
        )
        if use_prior:
            out_distr = self._distr_from_params(params_out)
            device = params_out[0].device
            prior_params = tuple(
                # use model output as "prior" for parameters not specified as prior
                (
                    torch.tensor(prior_params[i]).to(device)
                    if prior_params[i] is not None
                    else params_out[i]
                )
                for i in range(len(prior_params))
            )
            prior_distr = self._distr_from_params(prior_params)

            # Loss regularization using the prior distribution
            loss += self.prior_strength * torch.mean(
                kl_divergence(prior_distr, out_distr)
            )

        return loss

        Computes a loss from a `model_output`, which represents the parameters of a given probability
        distribution for every ground truth value in `target`, and the `target` itself.
        
darts/utils/likelihood_models/torch.py