From aedbd318d55681d23562d4a84e843aea9b0bd396 Mon Sep 17 00:00:00 2001
From: erogol <erogol@hotmail.com>
Date: Tue, 9 Jun 2020 23:02:28 +0200
Subject: [PATCH] remove weight norm for melgan generator

---
 vocoder/models/melgan_generator.py | 9 +++++++++
 1 file changed, 9 insertions(+)

diff --git a/vocoder/models/melgan_generator.py b/vocoder/models/melgan_generator.py
index 1e47816d..e69e6ef3 100644
--- a/vocoder/models/melgan_generator.py
+++ b/vocoder/models/melgan_generator.py
@@ -87,3 +87,12 @@ class MelganGenerator(nn.Module):
             (self.inference_padding, self.inference_padding),
             'replicate')
         return self.layers(cond_features)
+
+    def remove_weight_norm(self):
+        for _, layer in enumerate(self.layers):
+            if len(layer.state_dict()) != 0:
+                try:
+                    nn.utils.remove_weight_norm(layer)
+                except ValueError:
+                    layer.remove_weight_norm()
+