fix(ppo): identity function now takes 2 params, like message_passer

2026-05-07 14:07:19 +02:00 · 2026-05-07 14:07:19 +02:00 · 4e3cd0c673
commit 4e3cd0c673
parent 517c3134ca
1 changed files with 1 additions and 1 deletions
--- 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,