check for infs in inference

MIC-DKFZ · Dec 7, 2023 · eb3c40d · eb3c40d
1 parent 792343c
commit eb3c40d
Showing 1 changed file with 5 additions and 0 deletions.
diff --git a/nnunetv2/inference/predict_from_raw_data.py b/nnunetv2/inference/predict_from_raw_data.py
@@ -627,6 +627,11 @@ def predict_sliding_window_return_logits(self, input_image: torch.Tensor) \
                     n_predictions[sl[1:]] += (gaussian if self.use_gaussian else 1)
 
                 predicted_logits /= n_predictions
+                # check for infs
+                if torch.any(torch.isinf(predicted_logits)):
+                    raise RuntimeError('Encountered inf in predicted array. Aborting... If this problem persists, '
+                                       'reduce value_scaling_factor in compute_gaussian or increase the dtype of '
+                                       'predicted_logits to fp32')
         empty_cache(self.device)
         return predicted_logits[tuple([slice(None), *slicer_revert_padding[1:]])]