fix AdaLayerNorm (#94)

Author: akaitsuki-ii

diffsynth_engine/models/basic/transformer_helper.py

@@ -15,7 +15,7 @@ def __init__(self, dim: int, eps: float = 1e-6, device: str = "cuda:0", dtype: t
         self.silu = nn.SiLU()
 
     def forward(self, x, emb):
-        shift, scale = self.linear(self.silu(emb)).unsqueeze(1).chunk(2, dim=1)
+        shift, scale = self.linear(self.silu(emb)).unsqueeze(1).chunk(2, dim=2)
         return modulate(self.norm(x), shift, scale)
 
 
@@ -27,7 +27,7 @@ def __init__(self, dim, device: str, dtype: torch.dtype):
         self.norm = nn.LayerNorm(dim, elementwise_affine=False, eps=1e-6, device=device, dtype=dtype)
 
     def forward(self, x, emb):
-        shift, scale, gate = self.linear(self.silu(emb)).unsqueeze(1).chunk(3, dim=1)
+        shift, scale, gate = self.linear(self.silu(emb)).unsqueeze(1).chunk(3, dim=2)
         return modulate(self.norm(x), shift, scale), gate