fix tests

Author: Peter

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__':

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)