From f5645360dd1f45a7bf112f01c85fb862ee57df3c Mon Sep 17 00:00:00 2001
From: zhzluke96 <zhzluke96@outlook.com>
Date: Mon, 10 Jun 2024 20:59:34 +0800
Subject: [PATCH] :bug: fix hf.space force abort

---
 modules/Enhancer/ResembleEnhance.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/modules/Enhancer/ResembleEnhance.py b/modules/Enhancer/ResembleEnhance.py
index 553e454..a188a89 100644
--- 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()