423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-08-02 03:42:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

allow matrix_to_quaternion onnx export

Browse Source 
Summary: Attempt to allow torch.onnx.dynamo_export(matrix_to_quaternion) to work.

Differential Revision: D59812279

fbshipit-source-id: 4497e5b543bec9d5c2bdccfb779d154750a075ad
This commit is contained in:
Jeremy Reizenstein 2024-07-16 11:30:20 -07:00 committed by Facebook GitHub Bot
parent d0d0e02007
commit 7edaee71a9
1 changed files with 4 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
pytorch3d/transforms/rotation_conversions.py
View File

@ -97,7 +97,10 @@ def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
""" """
ret = torch.zeros_like(x) ret = torch.zeros_like(x)
positive_mask = x > 0 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 return ret