fix(rmvpe): pass `device` when loading torch model (#301)

2023-07-14 19:11:14 +08:00 · 2023-07-14 19:11:14 +08:00 · 55dd869858
parent e7b478596a
commit 55dd869858
1 changed files with 2 additions and 2 deletions
--- a/modules/F0Predictor/rmvpe/inference.py
+++ b/modules/F0Predictor/rmvpe/inference.py
@ -16,7 +16,7 @@ class RMVPE:
        else:
            self.device = device
        model = E2E0(4, 1, (2, 2))
-        ckpt = torch.load(model_path)
+        ckpt = torch.load(model_path, map_location=torch.device(self.device))
        model.load_state_dict(ckpt['model'])
        model = model.to(dtype).to(self.device)
        model.eval()