From aced74eeb559e23936038d36673003b0bd4110fd Mon Sep 17 00:00:00 2001 From: learned_optimization authors Date: Wed, 18 Mar 2026 11:40:17 -0700 Subject: [PATCH] No public description PiperOrigin-RevId: 885706162 --- learned_optimization/outer_trainers/gradient_learner.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/learned_optimization/outer_trainers/gradient_learner.py b/learned_optimization/outer_trainers/gradient_learner.py index a259aa3c..d6fb1749 100644 --- a/learned_optimization/outer_trainers/gradient_learner.py +++ b/learned_optimization/outer_trainers/gradient_learner.py @@ -428,6 +428,8 @@ def extract_one(idx, x): onp.asarray, estimator_out.unroll_info.task_param) iteration = estimator_out.unroll_info.iteration[ idx] if estimator_out.unroll_info.iteration is not None else None + if worker_weights.outer_state is None: + raise ValueError("worker_weights.outer_state is None") event_info.append({ "loss": estimator_out.unroll_info.loss[idx, :], "task_param": jax.tree_util.tree_map(fn, onp_task_params),