Add shape assertion to rotation code

cd41d73f · Gustaf Ahdritz · 253aff64 · cd41d73f
Commit cd41d73f authored May 07, 2022 by Gustaf Ahdritz
Show whitespace changes
Inline Side-by-side

Showing with 1 addition and 3 deletions

openfold/utils/rigid_utils.py openfold/utils/rigid_utils.py +1 -3

No files found.
--- a/openfold/utils/rigid_utils.py
+++ b/openfold/utils/rigid_utils.py
@@ -74,9 +74,7 @@ def rot_vec_mul(
        Returns:
            [*, 3] rotated coordinates
    """
-    x = t[..., 0]
+    x, y, z = torch.unbind(t, dim=-1)
-    y = t[..., 1]
-    z = t[..., 2]
    return torch.stack(
        [
            r[..., 0, 0] * x + r[..., 0, 1] * y + r[..., 0, 2] * z,