fix timing test

Author: Peter

src/pytorch_kinematics/chain.py

@@ -279,8 +279,8 @@ def get_link_names(self):
 
     @lru_cache
     def get_frame_indices(self, *frame_names):
-        return torch.tensor([self.frame_to_idx[n] for n in frame_names], dtype=torch.long,
-                            device=self.device)
+        return torch.tensor([self.frame_to_idx[n] for n in frame_names], dtype=torch.long, device=self.device)
+
     def forward_kinematics(self, th, frame_indices: Optional = None):
         if frame_indices is None:
             frame_indices = self.get_all_frame_indices()

tests/test_kinematics.py

@@ -115,6 +115,7 @@ def test_urdf_serial():
     print(chain)
     print(chain.get_joint_parameter_names())
     th = [0.0, -math.pi / 4.0, 0.0, math.pi / 2.0, 0.0, math.pi / 4.0, 0.0]
+
     ret = chain.forward_kinematics(th, end_only=False)
     tg = ret['lbr_iiwa_link_7']
     pos, rot = quat_pos_from_transform3d(tg)
@@ -126,8 +127,14 @@ def test_urdf_serial():
     dtype = torch.float64
 
     th_batch = torch.rand(N, len(chain.get_joint_parameter_names()), dtype=dtype, device=d)
+
     chain = chain.to(dtype=dtype, device=d)
 
+    # NOTE: Warmstart since pytorch can be slow the first time you run it
+    #  this has to be done after you move it to the GPU. Otherwise the timing isn't representative.
+    for _ in range(5):
+        ret = chain.forward_kinematics(th)
+
     number = 10
 
     def _fk_parallel():