remove weight norm for melgan generator

2020-06-09 23:02:28 +02:00 · 2020-06-09 23:02:28 +02:00 · c866f23af6
parent f9698af06e
commit c866f23af6
1 changed files with 9 additions and 0 deletions
--- 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()
+