diff --git a/sml/metrics/classification/classification_test.py b/sml/metrics/classification/classification_test.py
index 89c37bd0..cd4901d6 100644
--- a/sml/metrics/classification/classification_test.py
+++ b/sml/metrics/classification/classification_test.py
@@ -139,16 +139,14 @@ def check(spu_result, sk_result):
         # Test binary
         y_true = jnp.array([0, 1, 1, 0, 1, 1])
         y_pred = jnp.array([0, 0, 1, 0, 1, 1])
-        spu_result = spsim.sim_jax(sim, proc)(y_true, y_pred, pos_label=1, transform=0)
+        spu_result = spsim.sim_jax(sim, proc, static_argnums=(2, 5))(y_true, y_pred, 'binary', None, 1, False)
         sk_result = sklearn_proc(y_true, y_pred)
         check(spu_result, sk_result)
 
         # Test multiclass
         y_true = jnp.array([0, 1, 1, 0, 2, 1])
         y_pred = jnp.array([0, 0, 1, 0, 2, 1])
-        spu_result = spsim.sim_jax(sim, proc)(
-            y_true, y_pred, average=None, labels=[0, 1, 2]
-        )
+        spu_result = spsim.sim_jax(sim, proc, static_argnums=(2, 5))(y_true, y_pred, None, [0,1,2], 1, True)
         sk_result = sklearn_proc(y_true, y_pred, average=None, labels=[0, 1, 2])
         check(spu_result, sk_result)