fix: torch.load 添加 weights_only=False 以兼容新版 PyTorch
Co-Authored-By: Claude <noreply@anthropic.com>
This commit is contained in:
@@ -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()
|
||||||
|
|||||||
@@ -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()
|
||||||
|
|||||||
Reference in New Issue
Block a user