assertions

2026-04-23 16:54:56 +00:00 · 2026-02-20 02:16:49 -05:00 · 2026-02-20 02:16:49 -05:00 · e615eb1f50
commit e615eb1f50
parent 559d649a26
3 changed files with 86 additions and 23 deletions
--- a/atroposlib/envs/base.py
+++ b/atroposlib/envs/base.py
@ -1034,6 +1034,11 @@ class BaseEnv(ABC):
                            logger.error(f"[DISTILL] Failed to fetch teacher logprobs: {e}")
                            import traceback
                            logger.error(traceback.format_exc())
+                    self.teacher_client.assert_distill_arrays_aligned(
+                        token_sequences=group["tokens"],
+                        distill_token_ids=group.get("distill_token_ids"),
+                        distill_logprobs=group.get("distill_logprobs"),
+                    )
            else:
                logger.debug(
                    "[DISTILL] Skipped - enabled=%s, url=%s",