🐛 fix hf.space force abort

lenML · Jun 10, 2024 · f564536 · f564536
1 parent e969df8
commit f564536
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/modules/Enhancer/ResembleEnhance.py b/modules/Enhancer/ResembleEnhance.py
@@ -37,7 +37,7 @@ def load_model(self):
         enhancer = Enhancer(hparams)
         state_dict = torch.load(
             Path(MODELS_DIR) / "resemble-enhance" / "mp_rank_00_model_states.pt",
-            map_location="cpu",
+            map_location=self.device,
         )["module"]
         enhancer.load_state_dict(state_dict)
         enhancer.to(self.device).eval()