fix: torch.load 添加 weights_only=False 以兼容新版 PyTorch

Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
xsl
2026-07-09 23:44:48 +08:00
co-authored by Claude
parent 9675e147a0
commit 864d7f969a
2 changed files with 2 additions and 2 deletions
+1 -1
View File
@@ -79,7 +79,7 @@ class Resnet18(nn.Module):
def init_weight(self): def init_weight(self):
# 优先本地骨干权重(内网离线),缺失才回退 torch model_zoo(会查缓存)。 # 优先本地骨干权重(内网离线),缺失才回退 torch model_zoo(会查缓存)。
if os.path.isfile(_LOCAL_RESNET18): if os.path.isfile(_LOCAL_RESNET18):
state_dict = torch.load(_LOCAL_RESNET18, map_location="cpu") state_dict = torch.load(_LOCAL_RESNET18, map_location="cpu", weights_only=False)
else: else:
state_dict = modelzoo.load_url(resnet18_url) state_dict = modelzoo.load_url(resnet18_url)
self_state_dict = self.state_dict() self_state_dict = self.state_dict()
+1 -1
View File
@@ -52,7 +52,7 @@ class HairSegmenter:
self._torch = torch self._torch = torch
self.device = _select_device(torch) self.device = _select_device(torch)
self.net = BiSeNet(n_classes=N_CLASSES) self.net = BiSeNet(n_classes=N_CLASSES)
state = torch.load(weights_path, map_location="cpu") state = torch.load(weights_path, map_location="cpu", weights_only=False)
self.net.load_state_dict(state) self.net.load_state_dict(state)
self.net.to(self.device) self.net.to(self.device)
self.net.eval() self.net.eval()