diff --git a/src/MaxText/layers/nnx_wrappers.py b/src/MaxText/layers/nnx_wrappers.py
index 39baeaf985..1b74b2c531 100644
--- a/src/MaxText/layers/nnx_wrappers.py
+++ b/src/MaxText/layers/nnx_wrappers.py
@@ -29,7 +29,7 @@
 from flax.nnx import variablelib
 from flax.nnx.bridge import module as bdg_module
 from flax.nnx.module import Module
-from flax.nnx import Pytree
+from flax.nnx import Object # TODO(gulsumgudukbay) convert this import to from flax.nnx import Pytree after ROCm/jax supports flax >= 0.11.0
 from flax.nnx.rnglib import Rngs
 import jax
 from jax import tree_util as jtu
@@ -140,7 +140,7 @@ def nnx_attrs_to_linen_vars(nnx_attrs: dict) -> dict:
 
 def _set_initializing(module: Module, initializing: bool):
   for _, value in graph.iter_graph(module):
-    if isinstance(value, Pytree):
+    if isinstance(value, Object):
       value._object__state._initializing = initializing  # pylint: disable=protected-access