allow matrix_to_quaternion onnx export

Summary: Attempt to allow torch.onnx.dynamo_export(matrix_to_quaternion) to work. Differential Revision: D59812279 fbshipit-source-id: 4497e5b543bec9d5c2bdccfb779d154750a075ad
facebookresearch · Jul 16, 2024 · 7edaee7 · 7edaee7
1 parent d0d0e02
commit 7edaee7
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/pytorch3d/transforms/rotation_conversions.py b/pytorch3d/transforms/rotation_conversions.py
@@ -97,7 +97,10 @@ def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
     """
     ret = torch.zeros_like(x)
     positive_mask = x > 0
-    ret[positive_mask] = torch.sqrt(x[positive_mask])
+    if torch.is_grad_enabled():
+        ret[positive_mask] = torch.sqrt(x[positive_mask])
+    else:
+        ret = torch.where(positive_mask, torch.sqrt(x), ret)
     return ret