From db98fadcd065e03d35cde59d8b7691f82f166295 Mon Sep 17 00:00:00 2001 From: Linfeng Date: Mon, 9 Mar 2026 17:54:04 +0800 Subject: [PATCH] Fix checkpoint.load_config --- brax/training/checkpoint.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/brax/training/checkpoint.py b/brax/training/checkpoint.py index cc64d02a5..be332d250 100644 --- a/brax/training/checkpoint.py +++ b/brax/training/checkpoint.py @@ -225,6 +225,8 @@ def load_config( if init_fn_name not in loaded_dict['network_factory_kwargs']: continue init_fn_name_ = loaded_dict['network_factory_kwargs'][init_fn_name] + if init_fn_name_ is None: + continue loaded_dict['network_factory_kwargs'][init_fn_name] = ( networks.KERNEL_INITIALIZER[init_fn_name_] )