fix(PPOTrainer.py): bug fixes regarding jax.jit

2026-04-03 15:22:51 +02:00 · 2026-04-03 15:22:51 +02:00 · 62e11ce243
commit 62e11ce243
parent a9302cfd4e
2 changed files with 44 additions and 18 deletions
--- a/src/brittle_star_project/dataclasses/PPOArgs.py
+++ b/src/brittle_star_project/dataclasses/PPOArgs.py
@ -1,6 +1,9 @@
 from dataclasses import dataclass

+import jax

+
+@jax.tree_util.register_dataclass
@dataclass
 class PPOArgs:
    """