fix: release mixed-traffic warmup caches

1 files changed, 4 insertions, 0 deletions

diff --git a/experiments/conv_run.py b/experiments/conv_run.py
index d2c0653..2e2787b 100644
--- a/experiments/conv_run.py
+++ b/experiments/conv_run.py
@@ -386,6 +386,8 @@ def run(args):
             "uses_validation_endpoint": False,
         })
         counters["traffic_calibration_examples"] += count
+        del (calibration_forward, calibration_output_error,
+             calibration_instruction)
     if args.mode in ("wm", "rrm") and args.mirror_warmup_steps:
         sync(args.device)
         mirror_start = time.time()
@@ -430,6 +432,7 @@ def run(args):
                     config.nuisance_scale))
             predictor_metrics.append(predictor_metric)
             counters["predictor_warmup_examples"] += x.shape[0]
+        del forward
         train.g.set_state(loader_state)
         log["predictor_warmup"] = {
             "steps": args.predictor_warmup_steps,
@@ -447,6 +450,7 @@ def run(args):
                     net.predictor_traffic_residual_rms_ratio(
                         audit_forward["hiddens"]))
             counters["traffic_audit_examples"] += audit_x.shape[0]
+            del audit_forward
         sync(args.device)
         predictor_warmup_wall = time.time() - warmup_start