feat (stateful)(utilities): adding backward compatible tf_keras as …

…a supported base class.
ivy-llc · Sep 19, 2024 · 1fa2c53 · 1fa2c53
1 parent 7e8d08f
commit 1fa2c53
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/ivy/stateful/utilities.py b/ivy/stateful/utilities.py
@@ -15,7 +15,7 @@
 def _is_submodule(obj, kw):
     cls_str = {
         "torch": ("torch.nn.modules.module.Module",),
-        "keras": ("keras.engine.training.Model", "keras.src.models.model.Model", "keras.src.layers.layer.Layer"),
+        "keras": ("keras.engine.training.Model", "tf_keras.src.engine.training.Model", "keras.src.models.model.Model", "tf_keras.src.engine.base_layer.Layer", "keras.src.engine.base_layer.Layer", "keras.src.layers.layer.Layer"),
         "flax": ("flax.nnx.nnx.module.Module",),
     }[kw]
     try: