Added imports

prishajain1 · prishajain1 · commit 675554a44245 · 2025-12-17T10:34:04.000+05:30
diff --git a/src/maxdiffusion/models/embeddings_flax.py b/src/maxdiffusion/models/embeddings_flax.py
@@ -20,7 +20,7 @@
 import jax
 from .modeling_flax_utils import get_activation
 from ..models.attention_flax import NNXSimpleFeedForward
-from ..normalization_flax import FP32LayerNorm
+from ..models.normalization_flax import FP32LayerNorm
 
 
 def get_sinusoidal_embeddings(
