From 28e450fd5b9ee81cafb4d51f1deccd1c3be33df0 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <2225664821@qq.com>
Date: Sun, 20 Aug 2023 15:25:13 +0800
Subject: [PATCH] Update modules.py

---
 modules.py | 13 +++++++++++++
 1 file changed, 13 insertions(+)

diff --git a/modules.py b/modules.py
index 3f622f3..cc0aecf 100644
--- a/modules.py
+++ b/modules.py
@@ -16,7 +16,20 @@ from attentions import Encoder
 
 LRELU_SLOPE = 0.1
 
+class LayerNorm(nn.Module):
+  def __init__(self, channels, eps=1e-5):
+    super().__init__()
+    self.channels = channels
+    self.eps = eps
 
+    self.gamma = nn.Parameter(torch.ones(channels))
+    self.beta = nn.Parameter(torch.zeros(channels))
+
+  def forward(self, x):
+    x = x.transpose(1, -1)
+    x = F.layer_norm(x, (self.channels,), self.gamma, self.beta, self.eps)
+    return x.transpose(1, -1)
+    
 class ConvReluNorm(nn.Module):
   def __init__(self, in_channels, hidden_channels, out_channels, kernel_size, n_layers, p_dropout):
     super().__init__()