diff --git a/nerfstudio/field_components/spatial_distortions.py b/nerfstudio/field_components/spatial_distortions.py
index 449f4b5b9..02da548a3 100644
--- a/nerfstudio/field_components/spatial_distortions.py
+++ b/nerfstudio/field_components/spatial_distortions.py
@@ -65,12 +65,8 @@ class SceneContraction(SpatialDistortion):
 
     def forward(self, positions):
         def contract(x):
-            mag = torch.linalg.norm(x, ord=self.order, dim=-1)
-            mask = mag >= 1
-            x_new = x.clone()
-            x_new[mask] = (2 - (1 / mag[mask][..., None])) * (x[mask] / mag[mask][..., None])
-
-            return x_new
+            mag = torch.linalg.norm(x, ord=self.order, dim=-1)[..., None]
+            return torch.where(mag < 1, x, (2 - (1 / mag)) * (x / mag))
 
         if isinstance(positions, Gaussians):
             means = contract(positions.mean.clone())