Update FCPE.py

Bebra777228 · Aug 4, 2024 · 99e933c · 99e933c
1 parent 8256fe2
commit 99e933c
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/infer_pack/predictor/FCPE.py b/src/infer_pack/predictor/FCPE.py
@@ -664,7 +664,7 @@ def __init__(self, model_path, device=None, dtype=torch.float32):
         if device is None:
             device = "cuda" if torch.cuda.is_available() else "cpu"
         self.device = device
-        ckpt = torch.load(model_path, map_location=torch.device(self.device))
+        ckpt = torch.load(model_path, map_location=torch.device(self.device), weights_only=True)
         self.args = DotDict(ckpt["config"])
         self.dtype = dtype
         model = FCPE(