Merge pull request #119 from lincc-frameworks/add_loss_except

added exception to _get_loss function
grantmerz · Nov 20, 2024 · a516d3f · a516d3f
2 parents ac0176b + bfd97ad
commit a516d3f
Showing 1 changed file with 5 additions and 1 deletion.
diff --git a/src/deepdisc/astrodet/detectron.py b/src/deepdisc/astrodet/detectron.py
@@ -190,7 +190,11 @@ def _do_loss_eval(self):
 
     def _get_loss(self, data):
         # How loss is calculated on train_loop
-        metrics_dict = self._model(data)
+        try:
+            metrics_dict = self._model(data)
+        except:
+            print("Check the size of the images in the validation set")
+            return 0
         metrics_dict = {
             k: v.detach().cpu().item() if isinstance(v, torch.Tensor) else float(v)
             for k, v in metrics_dict.items()