move nan check to last diffusion step.

Author: janfb

sbi/samplers/score/diffuser.py

@@ -167,21 +167,20 @@ def run(
             # Apply predictor step
             samples = self.predictor(samples, t_current, t_next)
 
-            # Check for NaN values after predictor
-            if torch.isnan(samples).any():
-                raise RuntimeError(
-                    f"NaN values detected after predictor step "
-                    f"{time_step_idx}/{total_time_steps}. "
-                    f"This may indicate numerical instability in the vector field."
-                )
-
             # Apply corrector step if available
             if self.corrector is not None:
                 samples = self.corrector(samples, t_next, t_current)
 
             if save_intermediate:
                 intermediate_samples.append(samples)
 
+        # Check for NaN values after predictor
+        if torch.isnan(samples).any():
+            raise RuntimeError(
+                "NaN values detected after diffusion sampling "
+                "This may indicate numerical instability in the vector field."
+            )
+
         if save_intermediate:
             return torch.cat(intermediate_samples, dim=0)
         else: