add LayerNorm to ColorEncoder

piercus · Jan 22, 2024 · 37a21bf · 37a21bf
1 parent b1eaa66
commit 37a21bf
Showing 1 changed file with 8 additions and 1 deletion.
diff --git a/src/refiners/fluxion/adapters/color_palette.py b/src/refiners/fluxion/adapters/color_palette.py
@@ -55,10 +55,17 @@ def __init__(
         self,
         embedding_dim: int,
         device: Device | str | None = None,
+        eps: float = 1e-5,
         dtype: DType | None = None,
     ) -> None:
         super().__init__(
-            fl.Linear(in_features=4, out_features=embedding_dim, device=device, dtype=dtype),
+            fl.Linear(in_features=4, out_features=embedding_dim, bias=True, device=device, dtype=dtype),
+            fl.LayerNorm(
+                normalized_shape=embedding_dim,
+                eps=eps,
+                device=device,
+                dtype=dtype,
+            )
         )
 
 class ColorPaletteEncoder(fl.Chain):