From bc690632bbffee4d21dd504985e13bca2510bf13 Mon Sep 17 00:00:00 2001 From: Penn Date: Thu, 15 Dec 2022 13:24:31 -0800 Subject: [PATCH] Use `torch.where` for faster contract (#1132) use where for faster contract --- nerfstudio/field_components/spatial_distortions.py | 8 ++------ 1 file changed, 2 insertions(+), 6 deletions(-) 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())