Minor fix

NKI-AI · Aug 21, 2024 · d912b9e · d912b9e
1 parent f75f8e8
commit d912b9e
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/direct/nn/registration/registration.py b/direct/nn/registration/registration.py
@@ -114,7 +114,7 @@ def forward(self, moving_image: torch.Tensor, reference_image: torch.Tensor) ->
 
         # Estimate the displacement field
         displacement_field = [
-            self.displacement_transform(reference_image[_].cpu(), moving_image[_].cpu())
+            self.displacement_transform(reference_image[_].detach().cpu(), moving_image[_].detach().cpu())
             for _ in range(moving_image.shape[0])
         ]
         displacement_field = torch.stack(displacement_field, dim=0)