From a631235d469c919cac88c58c2b83e04c1834b03e Mon Sep 17 00:00:00 2001 From: YurenHao0426 Date: Wed, 22 Jul 2026 15:14:06 -0500 Subject: fix: release mixed-traffic warmup caches --- experiments/conv_run.py | 4 ++++ 1 file changed, 4 insertions(+) (limited to 'experiments') 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 -- cgit v1.2.3