From 5794eb54a99906eacef88765c9f431d9e4c4c089 Mon Sep 17 00:00:00 2001 From: Peter Mitrano Date: Fri, 3 Nov 2023 12:09:08 -0400 Subject: [PATCH] Update src/pytorch_kinematics/transforms/rotation_conversions.py Co-authored-by: Tom Power --- src/pytorch_kinematics/transforms/rotation_conversions.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/pytorch_kinematics/transforms/rotation_conversions.py b/src/pytorch_kinematics/transforms/rotation_conversions.py index dcc0ad4..bf75f98 100644 --- a/src/pytorch_kinematics/transforms/rotation_conversions.py +++ b/src/pytorch_kinematics/transforms/rotation_conversions.py @@ -497,8 +497,8 @@ def tensor_axis_and_angle_to_matrix(axis, theta): torch.stack([r10, r11, r12], -1), torch.stack([r20, r21, r22], -1)], -2) batch_shape = axis.shape[:-1] - mat44 = torch.eye(4, device=axis.device, dtype=axis.dtype).repeat(*batch_shape, 1, 1) - mat44[..., :3, :3] = rot + mat44 = torch.cat((rot, torch.zeros(*batch_shape, 3, 1).to(axis)), -1) + mat44 = torch.cat((mat44, torch.tensor([0.0, 0.0, 0.0, 1.0]).expand(*batch_shape, 1, 4).to(axis)), -2) return mat44