[fix] removing the loss_function stuff and replacing entirely with ev…

…al_step
uhlmanngroup · Aug 15, 2024 · fec2ee3 · fec2ee3
1 parent f67b815
commit fec2ee3
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/bioimage_embed/shapes/lightning.py b/bioimage_embed/shapes/lightning.py
@@ -31,7 +31,9 @@ def batch_to_tensor(self, batch):
 
         return output
 
-    def loss_function(self, model_output, *args, **kwargs):
+    def eval_step(self, batch, batch_idx):
+        # model_output = super().eval_step(batch, batch_idx)
+        model_output = self.predict_step(batch, batch_idx)
         loss_ops = lf.DistanceMatrixLoss(model_output.recon_x, norm=False)
 
         shape_loss = torch.sum(