diff --git a/src/brittle_star_project/ppo.py b/src/brittle_star_project/ppo.py
index fffc26b..fb27ccd 100644
--- a/src/brittle_star_project/ppo.py
+++ b/src/brittle_star_project/ppo.py
@@ -190,7 +190,7 @@ def ppo_loss(
     return loss, (pg_loss, v_loss, entropy_loss, jax.lax.stop_gradient(approx_kl))
 
 
-def identity(hidden):
+def identity(_, hidden):
     """
     Used for seamless jax integration,
     avoids having branching inside jitted function,