Don't crash with ZeroDivisionError if initial test ELBO is the best t…

…est ELBO
broadinstitute · Apr 22, 2024 · 4e8ffb9 · 4e8ffb9
1 parent 72e6a3f
commit 4e8ffb9
Showing 1 changed file with 6 additions and 1 deletion.
diff --git a/cellbender/remove_background/train.py b/cellbender/remove_background/train.py
@@ -247,7 +247,12 @@ def run_training(model: RemoveBackgroundPyroModel,
             if model.loss['test']['elbo'][-1] < best_test_elbo:
                 final_best_diff = best_test_elbo - model.loss['test']['elbo'][-1]
                 initial_best_diff = best_test_elbo - model.loss['test']['elbo'][0]
-                if (final_best_diff / initial_best_diff) > final_elbo_fail_fraction:
+                if initial_best_diff == 0:
+                    raise ElboException(
+                        f"Training failed because there was no improvement from the initial test loss {model.loss['test']['elbo'][0]:.2f}. "
+                        f"Final test loss was {model.loss['test']['elbo'][-1]}"
+                    )
+                elif (final_best_diff / initial_best_diff) > final_elbo_fail_fraction:
                     raise ElboException(
                         f"Training failed because final test loss {model.loss['test']['elbo'][-1]:.2f} "
                         f'is not sufficiently close to best test loss {best_test_elbo:.2f}, '