diff --git a/etils/ecolab/lazy_imports.py b/etils/ecolab/lazy_imports.py
index e8338861..a13d1c8a 100644
--- a/etils/ecolab/lazy_imports.py
+++ b/etils/ecolab/lazy_imports.py
@@ -154,6 +154,7 @@ def print_current_imports() -> None:
   import flask
   import flax
   from flax import linen as nn
+  from flax import nnx
   import functorch
   import gin
   import grain.python as grain