From 1a223661ab01bcee0effe73839c61122518ce4fa Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Stardust=C2=B7=E5=87=8F?= <2225664821@qq.com>
Date: Mon, 21 Aug 2023 08:32:39 +0800
Subject: [PATCH] Update models.py

---
 models.py | 8 +++++---
 1 file changed, 5 insertions(+), 3 deletions(-)

diff --git a/models.py b/models.py
index a8660e7..a714bac 100644
--- a/models.py
+++ b/models.py
@@ -182,7 +182,8 @@ class TextEncoder(nn.Module):
                  n_heads,
                  n_layers,
                  kernel_size,
-                 p_dropout):
+                 p_dropout,
+                 gin_channels=0):
         super().__init__()
         self.n_vocab = n_vocab
         self.out_channels = out_channels
@@ -192,7 +193,7 @@ class TextEncoder(nn.Module):
         self.n_layers = n_layers
         self.kernel_size = kernel_size
         self.p_dropout = p_dropout
-
+        self.gin_channels = gin_channels
         self.emb = nn.Embedding(len(symbols), hidden_channels)
         nn.init.normal_(self.emb.weight, 0.0, hidden_channels ** -0.5)
         self.tone_emb = nn.Embedding(num_tones, hidden_channels)
@@ -207,7 +208,8 @@ class TextEncoder(nn.Module):
             n_heads,
             n_layers,
             kernel_size,
-            p_dropout)
+            p_dropout,
+            gin_channels=self.gin_channels)
         self.proj = nn.Conv1d(hidden_channels, out_channels * 2, 1)
 
     def forward(self, x, x_lengths, tone, language, bert):