fix: add nan check to _loss methods. (#1361)

* add nan check to _loss method.

* add checks for nres

Author: janfb

sbi/inference/trainers/nle/nle_base.py

@@ -22,6 +22,7 @@
     reshape_to_batch_event,
 )
 from sbi.utils import check_estimator_arg, check_prior, x_shape_from_simulation
+from sbi.utils.torchutils import assert_all_finite
 
 
 class LikelihoodEstimator(NeuralInference, ABC):
@@ -381,4 +382,6 @@ def _loss(self, theta: Tensor, x: Tensor) -> Tensor:
             theta, event_shape=self._neural_net.condition_shape
         )
         x = reshape_to_batch_event(x, event_shape=self._neural_net.input_shape)
-        return self._neural_net.loss(x, condition=theta)
+        loss = self._neural_net.loss(x, condition=theta)
+        assert_all_finite(loss, "NLE loss")
+        return loss

sbi/inference/trainers/npe/npe_base.py

@@ -42,6 +42,7 @@
     warn_if_zscoring_changes_data,
 )
 from sbi.utils.sbiutils import ImproperEmpirical, mask_sims_from_prior
+from sbi.utils.torchutils import assert_all_finite
 
 
 class PosteriorEstimator(NeuralInference, ABC):
@@ -609,6 +610,7 @@ def _loss(
             # Must be extended ones other Estimators are implemented. See #966,
             loss = -self._log_prob_proposal_posterior(theta, x, masks, proposal)
 
+        assert_all_finite(loss, "NPE loss")
         return calibration_kernel(x) * loss
 
     def _check_proposal(self, proposal):

sbi/inference/trainers/npse/npse.py

@@ -28,6 +28,7 @@
     warn_if_zscoring_changes_data,
 )
 from sbi.utils.sbiutils import ImproperEmpirical, mask_sims_from_prior
+from sbi.utils.torchutils import assert_all_finite
 
 
 class NPSE(NeuralInference):
@@ -510,6 +511,7 @@ def _loss(
                 "Multi-round NPSE with arbitrary proposals is not implemented"
             )
 
+        assert_all_finite(loss, "NPSE loss")
         return calibration_kernel(x) * loss
 
     def _converged(self, epoch: int, stop_after_epochs: int) -> bool:

sbi/inference/trainers/nre/bnre.py

@@ -10,6 +10,7 @@
 from sbi.inference.trainers.nre.nre_a import NRE_A
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
+from sbi.utils.torchutils import assert_all_finite
 
 
 class BNRE(NRE_A):
@@ -142,4 +143,6 @@ def _loss(
             .square()
         )
 
-        return bce + regularization_strength * regularizer
+        loss = bce + regularization_strength * regularizer
+        assert_all_finite(loss, "BNRE loss")
+        return loss

sbi/inference/trainers/nre/nre_a.py

@@ -10,6 +10,7 @@
 from sbi.inference.trainers.nre.nre_base import RatioEstimator
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
+from sbi.utils.torchutils import assert_all_finite
 
 
 class NRE_A(RatioEstimator):
@@ -124,4 +125,6 @@ def _loss(self, theta: Tensor, x: Tensor, num_atoms: int) -> Tensor:
         labels[1::2] = 0.0
 
         # Binary cross entropy to learn the likelihood (AALR-specific)
-        return nn.BCELoss()(likelihood, labels)
+        loss = nn.BCELoss()(likelihood, labels)
+        assert_all_finite(loss, "NRE-A loss")
+        return loss

sbi/inference/trainers/nre/nre_b.py

@@ -10,6 +10,7 @@
 from sbi.inference.trainers.nre.nre_base import RatioEstimator
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
+from sbi.utils.torchutils import assert_all_finite
 
 
 class NRE_B(RatioEstimator):
@@ -122,4 +123,6 @@ def _loss(self, theta: Tensor, x: Tensor, num_atoms: int) -> Tensor:
         # "correct" one for the 1-out-of-N classification.
         log_prob = logits[:, 0] - torch.logsumexp(logits, dim=-1)
 
-        return -torch.mean(log_prob)
+        loss = -torch.mean(log_prob)
+        assert_all_finite(loss, "NRE-B loss")
+        return loss

sbi/inference/trainers/nre/nre_c.py

@@ -10,6 +10,7 @@
 from sbi.inference.trainers.nre.nre_base import RatioEstimator
 from sbi.sbi_types import TensorboardSummaryWriter
 from sbi.utils.sbiutils import del_entries
+from sbi.utils.torchutils import assert_all_finite
 
 
 class NRE_C(RatioEstimator):
@@ -190,7 +191,10 @@ def _loss(
 
         # relative weights. p_marginal := p_0, and p_joint := p_K * K from the notation.
         p_marginal, p_joint = self._get_prior_probs_marginal_and_joint(gamma)
-        return -torch.mean(p_marginal * log_prob_marginal + p_joint * log_prob_joint)
+
+        loss = -torch.mean(p_marginal * log_prob_marginal + p_joint * log_prob_joint)
+        assert_all_finite(loss, "NRE-C loss")
+        return loss
 
     @staticmethod
     def _get_prior_probs_marginal_and_joint(gamma: float) -> Tuple[float, float]: