Make dx, dy, dz consistent (eliminates broadcasting)

uncbiag · Apr 14, 2024 · 4beb7e3 · 4beb7e3
1 parent 6c13d48
commit 4beb7e3
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/icon_registration/losses.py b/src/icon_registration/losses.py
@@ -153,7 +153,7 @@ def compute_gradient_icon_loss(self, phi_AB, phi_BA):
             dy = torch.Tensor([[[[[0.0]]], [[[delta]]], [[[0.0]]]]]).to(
                 self.identity_map.device
             )
-            dz = torch.Tensor([[[[0.0]]], [[[0.0]]], [[[delta]]]]).to(
+            dz = torch.Tensor([[[[[0.0]]], [[[0.0]]], [[[delta]]]]]).to(
                 self.identity_map.device
             )
             direction_vectors = (dx, dy, dz)