[feat] adding model_predict_test

uhlmanngroup · Aug 15, 2024 · 1552a87 · 1552a87
1 parent 16f431d
commit 1552a87
Showing 1 changed file with 6 additions and 0 deletions.
diff --git a/bioimage_embed/shapes/tests/test_lightning.py b/bioimage_embed/shapes/tests/test_lightning.py
@@ -79,3 +79,9 @@ def test_model(trainer, lit_model, dataloader):
 
 def test_model_fit(trainer, lit_model, dataloader):
     return trainer.fit(lit_model, dataloader)
+
+
+def test_model_predict(trainer, lit_model, dataloader):
+    y = trainer.predict(lit_model, dataloader)
+    # TODO Add checks for shape_loss and potentially other losses (contrastive loss)
+    return y