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
+2 -2
View File
@@ -216,9 +216,9 @@ def is_url(url_or_filename):
def load_checkpoint(model,url_or_filename):
if is_url(url_or_filename):
cached_file = download_cached_file(url_or_filename, check_hash=False, progress=True)
checkpoint = torch.load(cached_file, map_location='cpu')
checkpoint = torch.load(cached_file, map_location='cpu', weights_only=True)
elif os.path.isfile(url_or_filename):
checkpoint = torch.load(url_or_filename, map_location='cpu')
checkpoint = torch.load(url_or_filename, map_location='cpu', weights_only=True)
else:
raise RuntimeError('checkpoint url or path is invalid')