audit: enforce paired task order after neutral warmup

1 files changed, 4 insertions, 0 deletions

diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index c25c48c..a1ebb6a 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -439,12 +439,16 @@ def run(args):
             counters["predictor_warmup_examples"] += x.shape[0]
         del forward
         train.g.set_state(loader_state)
+        loader_state_restored = torch.equal(train.g.get_state(), loader_state)
+        if not loader_state_restored:
+            raise AssertionError("predictor warmup did not restore loader state")
         log["predictor_warmup"] = {
             "steps": args.predictor_warmup_steps,
             "first_mse": predictor_metrics[0],
             "mean_mse": sum(predictor_metrics) / len(predictor_metrics),
             "last_mse": predictor_metrics[-1],
             "instruction_present": False,
+            "task_loader_state_restored": loader_state_restored,
         }
         if args.mode == "kp_traffic":
             audit_x, _ = traffic_calibration_batch