diff --git a/mjinx/components/tasks/_com_task.py b/mjinx/components/tasks/_com_task.py
index 4c47c96..8a4835c 100644
--- a/mjinx/components/tasks/_com_task.py
+++ b/mjinx/components/tasks/_com_task.py
@@ -137,7 +137,7 @@ def target_com(self) -> jnp.ndarray:
         return self._target_com
 
     @target_com.setter
-    def target_com(self, value: Sequence):
+    def target_com(self, value: jnp.ndarray | Sequence):
         """
         Set the target center of mass for the task.
 
@@ -145,7 +145,7 @@ def target_com(self, value: Sequence):
         """
         self.update_target_com(value)
 
-    def update_target_com(self, target_com: Sequence):
+    def update_target_com(self, target_com: jnp.ndarray | Sequence):
         """
         Update the target center of mass for the task.
 
diff --git a/mjinx/components/tasks/_joint_task.py b/mjinx/components/tasks/_joint_task.py
index 87a4034..707e843 100644
--- a/mjinx/components/tasks/_joint_task.py
+++ b/mjinx/components/tasks/_joint_task.py
@@ -177,7 +177,7 @@ def target_q(self) -> jnp.ndarray:
         return self._target_q
 
     @target_q.setter
-    def target_q(self, value: Sequence):
+    def target_q(self, value: jnp.ndarray | Sequence):
         """
         Set the target joint positions for the task.
 
@@ -185,7 +185,7 @@ def target_q(self, value: Sequence):
         """
         self.update_target_q(value)
 
-    def update_target_q(self, target_q: Sequence):
+    def update_target_q(self, target_q: jnp.ndarray | Sequence):
         """
         Update the target joint positions for the task.
 
diff --git a/mjinx/typing.py b/mjinx/typing.py
index f97e1ec..7ce3083 100644
--- a/mjinx/typing.py
+++ b/mjinx/typing.py
@@ -11,6 +11,7 @@
 import jax
 import jax.numpy as jnp
 import numpy as np
+import numpy.typing
 from mujoco.mjx._src.dataclasses import PyTreeNode
 
 ndarray: TypeAlias = np.ndarray | jnp.ndarray
diff --git a/pyproject.toml b/pyproject.toml
index 7b96daa..2c1614e 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -57,4 +57,4 @@ line-length = 120
 "__init__.py" = ["F401"]
 
 [tool.pytest]
-filterwarnings = "ignore:.*U.*mode is deprecated:DeprecationWarning"
+filterwarnings = ["ignore:.*U.*mode is deprecated:DeprecationWarning"]
diff --git a/tests/unit_tests/configuration/test_configuration.py b/tests/unit_tests/configuration/test_configuration.py
index efdbece..7100008 100644
--- a/tests/unit_tests/configuration/test_configuration.py
+++ b/tests/unit_tests/configuration/test_configuration.py
@@ -305,8 +305,9 @@ def test_get_distance_with_hfield(self) -> None:
        </mujoco>
        """
         )
-        hfield_data = mjx.make_data(hfield_model)
-        hfield_data = mjx.kinematics(hfield_model, hfield_data)
+        hfield_model_mjx = mjx.put_model(hfield_model)
+        hfield_data = mjx.make_data(hfield_model_mjx)
+        hfield_data = mjx.kinematics(hfield_model_mjx, hfield_data)
 
         collision_pairs: list[CollisionPair] = [(0, 1)]  # hfield and sphere