Expanding odak.learn.wave.multiplane_loss().

kaanaksit · Jul 10, 2024 · 5aaf37b · 5aaf37b
1 parent 93409a1
commit 5aaf37b
Showing 1 changed file with 4 additions and 3 deletions.
diff --git a/odak/learn/wave/loss.py b/odak/learn/wave/loss.py
@@ -310,9 +310,10 @@ def __call__(self, image, target, plane_id = None):
         l2_cor = self.weights[2] * self.loss_function(image * target, target * target)
         loss = l2 + l2_mask + l2_cor
         try:
-            if self.weights[3] > 0.:
-                l_ColorVideoVDP = self.weights[3] * self.cvvdp.loss(image.unsqueeze(0), target.unsqueeze(0), dim_order = 'CHW')
-                loss += l_ColorVideoVDP
+            if len(self.weights) > 3:
+                if self.weights[3] > 0.:
+                    l_ColorVideoVDP = self.weights[3] * self.cvvdp.loss(image.unsqueeze(0), target.unsqueeze(0), dim_order = 'CHW')
+                    loss += l_ColorVideoVDP
         except:
             logging.warning('ColorVideoVDP loss failed for some reason. Could ColorVideoVDP missing? If so, visit GitHub:gfxdisp/ColorVideoVDP')
         return loss