From 241e968df17eed026113c9e17e2893a103f09d23 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Eren=20G=C3=B6lge?= Date: Tue, 6 Apr 2021 11:02:50 +0200 Subject: [PATCH] load_checkpoint for hifigan and no_grad for inference --- TTS/vocoder/models/hifigan_generator.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/TTS/vocoder/models/hifigan_generator.py b/TTS/vocoder/models/hifigan_generator.py index 44b0edfa..1da726c6 100644 --- a/TTS/vocoder/models/hifigan_generator.py +++ b/TTS/vocoder/models/hifigan_generator.py @@ -159,6 +159,7 @@ class HifiganGenerator(torch.nn.Module): x = torch.tanh(x) return x + @torch.no_grad() def inference(self, c): c = c.to(self.conv_pre.weight.device) c = torch.nn.functional.pad( @@ -173,3 +174,11 @@ class HifiganGenerator(torch.nn.Module): l.remove_weight_norm() remove_weight_norm(self.conv_pre) remove_weight_norm(self.conv_post) + + def load_checkpoint(self, config, checkpoint_path, eval=False): # pylint: disable=unused-argument, redefined-builtin + state = torch.load(checkpoint_path, map_location=torch.device('cpu')) + self.load_state_dict(state['model']) + if eval: + self.eval() + assert not self.training + self.remove_weight_norm()