From 614ca1b142d8c361ce29f7097e4091f9a0a44c43 Mon Sep 17 00:00:00 2001 From: Alireza Omidi Date: Sat, 3 Oct 2020 18:35:56 +0330 Subject: [PATCH] Resolved issue in loading the pretrained model for CPU --- hubconf.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/hubconf.py b/hubconf.py index 1031f6b..cd42923 100644 --- a/hubconf.py +++ b/hubconf.py @@ -10,13 +10,14 @@ } -def melgan(model_name='nvidia_tacotron2_LJ11_epoch6400', pretrained=True, progress=True): +def melgan(model_name='nvidia_tacotron2_LJ11_epoch6400', pretrained=True, progress=True, device='cuda'): params = model_params[model_name] model = Generator(params['mel_channel']) if pretrained: state_dict = torch.hub.load_state_dict_from_url(params['model_url'], - progress=progress) + progress=progress, + map_location=torch.device(device)) model.load_state_dict(state_dict['model_g']) model.eval(inference=True)