Delete layer checkpoint deps

nkovela1 · Dec 14, 2023 · d83362d · d83362d
1 parent e669c27
commit d83362d
Showing 1 changed file with 8 additions and 6 deletions.
diff --git a/keras_nlp/utils/preset_utils.py b/keras_nlp/utils/preset_utils.py
@@ -165,12 +165,14 @@ def legacy_load_weights(layer, weights_path):
     # We find the `Functional` class, and temporarily remove the
     # `_layer_checkpoint_dependencies` property, which on older version of
     # TensorFlow complete broke the variable paths for functional models.
-    property = layer._layer_checkpoint_dependencies
-    layer._layer_checkpoint_dependencies = None
-    try:
-        layer.load_weights(weights_path)
-    finally:
-        layer._layer_checkpoint_dependencies = property
+    functional_cls = None
+    for cls in inspect.getmro(layer.__class__):
+        if cls.__name__ == "Functional":
+            functional_cls = cls
+    property = functional_cls._layer_checkpoint_dependencies
+    del functional_cls._layer_checkpoint_dependencies
+    layer.load_weights(weights_path)
+    functional_cls._layer_checkpoint_dependencies = property
 
 
 def load_from_preset(