update for SpaSEG

STOmics · Aug 19, 2024 · 2dcaff1 · 2dcaff1
1 parent d9d9faa
commit 2dcaff1
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/stereo/algorithm/spa_seg/spaseg.py b/stereo/algorithm/spa_seg/spaseg.py
@@ -115,8 +115,8 @@ def _train(self,
             loss_MSE = torch.nn.MSELoss()
 
             # edge loss definition
-            loss_edge_vertical = torch.nn.L1Loss(size_average=True)
-            loss_edge_horizontal = torch.nn.L1Loss(size_average=True)
+            loss_edge_vertical = torch.nn.L1Loss(reduction='mean')
+            loss_edge_horizontal = torch.nn.L1Loss(reduction='mean')
 
             vertical_target = torch.zeros(data.shape[0], im.shape[1] - 1, im.shape[2], self.output_dim)
             horizontal_target = torch.zeros(data.shape[0], im.shape[1], im.shape[2] - 1, self.output_dim)