Deactivate failsafe in hk.flax.lift to avoid false-positives.

PiperOrigin-RevId: 607348942
google-deepmind · Feb 15, 2024 · 63c4a6f · 63c4a6f
1 parent 134763d
commit 63c4a6f
Showing 1 changed file with 1 addition and 0 deletions.
diff --git a/haiku/_src/flax/transform_flax.py b/haiku/_src/flax/transform_flax.py
@@ -185,6 +185,7 @@ def lift(
   """
   mod = _flax_transform_with_state(mod)
   init_fn, updater = hk.lift_with_state(mod.init, name=name, allow_reuse=True)
+  updater._used = True  # pylint: disable=protected-access
 
   def wrapped(*args, **kwargs):
     init_kwargs = dict(kwargs)