From 5aaf37bdcaf32dbe856aef797399ad191a2a37be Mon Sep 17 00:00:00 2001 From: kaanaksit Date: Wed, 10 Jul 2024 19:31:11 +0100 Subject: [PATCH] Expanding odak.learn.wave.multiplane_loss(). --- odak/learn/wave/loss.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/odak/learn/wave/loss.py b/odak/learn/wave/loss.py index f57ca979..e9fdf854 100644 --- 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