fix tests

UM-ARM-Lab · Oct 30, 2023 · 4385a64 · 4385a64
1 parent e2909b6
commit 4385a64
Show file tree

Hide file tree

Showing 3 changed files with 29 additions and 24 deletions.
diff --git a/tests/test_menagerie.py b/tests/test_menagerie.py
@@ -44,7 +44,7 @@ def test_menagerie():
         print(f"\t {chain.get_frame_names()}")
         print(f"\t {chain.get_joint_parameter_names()}")
         th = np.zeros(len(chain.get_joint_parameter_names()))
-        fk_dict = chain.forward_kinematics(th, end_only=True)
+        fk_dict = chain.forward_kinematics(th)
 
 
 if __name__ == '__main__':

diff --git a/tests/test_rotation_conversions.py b/tests/test_rotation_conversions.py
@@ -11,8 +11,9 @@ def test_axis_angle_to_matrix_perf():
     number = 100
     N = 1_000
 
-    axis_angle = torch.randn([N, 3], device='cuda', dtype=torch.float64)
-    axis_1d = torch.tensor([1., 0, 0], device='cuda', dtype=torch.float64)  # in the FK code this is NOT batched!
+    device = 'cuda' if torch.cuda.is_available() else 'cpu'
+    axis_angle = torch.randn([N, 3], device=device, dtype=torch.float64)
+    axis_1d = torch.tensor([1., 0, 0], device=device, dtype=torch.float64)  # in the FK code this is NOT batched!
     theta = axis_angle.norm(dim=1, keepdim=True)
 
     dt1 = timeit.timeit(lambda: axis_angle_to_matrix(axis_angle), number=number)