Use weights_only for loading (#3427)

Co-authored-by: Manuel Schmid <9307310+mashb1t@users.noreply.github.com>
This commit is contained in:
Sergii Dymchenko
2024-08-03 12:33:01 +02:00
committed by GitHub
co-authored by Manuel Schmid
parent 1a53e0676a
commit da3d4d006f
14 changed files with 21 additions and 21 deletions
+1 -1
View File
@@ -17,7 +17,7 @@ def init_parsing_model(model_name='bisenet', half=False, device='cuda', model_ro
model_path = load_file_from_url(
url=model_url, model_dir='facexlib/weights', progress=True, file_name=None, save_dir=model_rootpath)
load_net = torch.load(model_path, map_location=lambda storage, loc: storage)
load_net = torch.load(model_path, map_location=lambda storage, loc: storage, weights_only=True)
model.load_state_dict(load_net, strict=True)
model.eval()
model = model.to(device)