From 2b97de706ca136bb19026d0b8b2981ed5578c92e Mon Sep 17 00:00:00 2001
From: Haoming <hmstudy02@gmail.com>
Date: Thu, 19 Feb 2026 11:34:07 +0800
Subject: [PATCH] chroma

---
 backend/nn/chroma.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/backend/nn/chroma.py b/backend/nn/chroma.py
index fbf6bb89..ba8b66f7 100644
--- a/backend/nn/chroma.py
+++ b/backend/nn/chroma.py
@@ -11,7 +11,6 @@ from backend.nn.flux import (
     EmbedND,
     MLPEmbedder,
     ModulationOut,
-    RMSNorm,
     SingleStreamBlock,
     timestep_embedding,
 )
@@ -33,7 +32,7 @@ class Approximator(nn.Module):
         super().__init__()
         self.in_proj = nn.Linear(in_dim, hidden_dim, bias=True)
         self.layers = nn.ModuleList([MLPEmbedder(hidden_dim, hidden_dim) for _ in range(n_layers)])
-        self.norms = nn.ModuleList([RMSNorm(hidden_dim) for _ in range(n_layers)])
+        self.norms = nn.ModuleList([nn.RMSNorm(hidden_dim) for _ in range(n_layers)])
         self.out_proj = nn.Linear(hidden_dim, out_dim)
 
     @property