diff --git a/embodichain/lab/sim/objects/rigid_object.py b/embodichain/lab/sim/objects/rigid_object.py index a6b650c..c701236 100644 --- a/embodichain/lab/sim/objects/rigid_object.py +++ b/embodichain/lab/sim/objects/rigid_object.py @@ -350,6 +350,7 @@ def set_local_pose( # we should keep `pose_` life cycle to the end of the function. pose = torch.cat((quat, xyz), dim=-1) indices = self.body_data.gpu_indices[local_env_ids] + torch.cuda.synchronize(self.device) self._ps.gpu_apply_rigid_body_data( data=pose.clone(), gpu_indices=indices, @@ -458,6 +459,7 @@ def add_force_torque( else: indices = self.body_data.gpu_indices[local_env_ids] + torch.cuda.synchronize(self.device) if force is not None: self._ps.gpu_apply_rigid_body_data( data=force, @@ -658,6 +660,7 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: (len(local_env_ids), 3), dtype=torch.float32, device=self.device ) indices = self.body_data.gpu_indices[local_env_ids] + torch.cuda.synchronize(self.device) self._ps.gpu_apply_rigid_body_data( data=zeros, gpu_indices=indices, diff --git a/embodichain/lab/sim/objects/rigid_object_group.py b/embodichain/lab/sim/objects/rigid_object_group.py index baf62f4..f22a245 100644 --- a/embodichain/lab/sim/objects/rigid_object_group.py +++ b/embodichain/lab/sim/objects/rigid_object_group.py @@ -362,6 +362,7 @@ def set_local_pose( indices = self.body_data.gpu_indices[local_env_ids][ :, local_obj_ids ].flatten() + torch.cuda.synchronize(self.device) self._ps.gpu_apply_rigid_body_data( data=pose.clone(), gpu_indices=indices, @@ -433,6 +434,7 @@ def clear_dynamics(self, env_ids: Sequence[int] | None = None) -> None: device=self.device, ) indices = self.body_data.gpu_indices[local_env_ids].flatten() + torch.cuda.synchronize(self.device) self._ps.gpu_apply_rigid_body_data( data=zeros, gpu_indices=indices,