fix: renaming the USE_NATIVE_KERAS_LAYERS env variable

ivy-llc · Sep 18, 2024 · 27dc926 · 27dc926
1 parent 47ca552
commit 27dc926
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/ivy/stateful/utilities.py b/ivy/stateful/utilities.py
@@ -89,9 +89,9 @@ def _maybe_update_flax_layer_weights(layer, weight_name, new_weight):
     import flax.nnx as nnx
     import jax.numpy as jnp
 
-    has_keras_layers = os.environ.get("USE_NATIVE_KERAS_LAYERS", None) == "true"
+    has_flax_layers = os.environ.get("USE_NATIVE_FW_LAYERS", None) == "true"
     transpose_weights = (
-        has_keras_layers
+        has_flax_layers
         or os.environ.get("APPLY_TRANSPOSE_OPTIMIZATION", None) == "true"
     )
 
@@ -286,7 +286,7 @@ def _maybe_update_keras_layer_weights(layer, weight_name, new_weight):
     else:
         KerasVariable = tf.Variable
 
-    has_keras_layers = os.environ.get("USE_NATIVE_KERAS_LAYERS", None) == "true"
+    has_keras_layers = os.environ.get("USE_NATIVE_FW_LAYERS", None) == "true"
     transpose_weights = (
         has_keras_layers
         or os.environ.get("APPLY_TRANSPOSE_OPTIMIZATION", None) == "true"