Update losses.py

uncbiag · May 24, 2024 · db6cd42 · db6cd42
1 parent f9d6b6b
commit db6cd42
Showing 1 changed file with 4 additions and 4 deletions.
diff --git a/src/icon_registration/losses.py b/src/icon_registration/losses.py
@@ -43,7 +43,7 @@ def forward(self, image_A, image_B):
 
         # Tag used elsewhere for optimization.
         # Must be set at beginning of forward b/c not preserved by .cuda() etc
-        self.tag_identity()
+        self.tag_identity_map()
 
         self.phi_AB = self.regis_net(image_A, image_B)
         self.phi_BA = self.regis_net(image_B, image_A)
@@ -209,7 +209,7 @@ def forward(self, image_A, image_B) -> ICONLoss:
 
         # Tag used elsewhere for optimization.
         # Must be set at beginning of forward b/c not preserved by .cuda() etc
-        self.tag_identity()
+        self.tag_identity_map()
 
         self.phi_AB = self.regis_net(image_A, image_B)
         self.phi_BA = self.regis_net(image_B, image_A)
@@ -252,7 +252,7 @@ def forward(self, image_A, image_B):
 
         # Tag used elsewhere for optimization.
         # Must be set at beginning of forward b/c not preserved by .cuda() etc
-        self.tag_identity()
+        self.tag_identity_map()
 
         self.phi_AB = self.regis_net(image_A, image_B)
         self.phi_BA = self.regis_net(image_B, image_A)
@@ -454,7 +454,7 @@ def forward(self, image_A, image_B) -> ICONLoss:
 
         # Tag used elsewhere for optimization.
         # Must be set at beginning of forward b/c not preserved by .cuda() etc
-        self.identity_map.isIdentity = True
+        self.tag_identity_map()
 
         self.phi_AB = self.regis_net(image_A, image_B)
         self.phi_AB_vectorfield = self.phi_AB(self.identity_map)