Merge pull request #67 from uncbiag/nanfix

Update losses.py
uncbiag · Feb 3, 2024 · f23a698 · f23a698
2 parents 50d2cba + 8dee552
commit f23a698
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/src/icon_registration/losses.py b/src/icon_registration/losses.py
@@ -592,8 +592,8 @@ def __call__(self, image_A, image_B):
             1
             - (self.blur(I * J) - (self.blur(I) * self.blur(J)))
             / torch.sqrt(
-                (self.blur(I * I) - self.blur(I) ** 2 + 0.00001)
-                * (self.blur(J * J) - self.blur(J) ** 2 + 0.00001)
+                (torch.relu(self.blur(I * I) - self.blur(I) ** 2) + 0.00001)
+                * (torch.relu(self.blur(J * J) - self.blur(J) ** 2) + 0.00001)
             )
         )