"git@developer.sourcefind.cn:wangsen/paddle_dbnet.git" did not exist on "f8f715352688656dcf0e6f3574bf7d988a1142b5"
Commit e1278fa9 authored by comfyanonymous's avatar comfyanonymous
Browse files

Support old pytorch versions that don't have weights_only.

parent 9b1396e9
...@@ -6,6 +6,10 @@ def load_torch_file(ckpt, safe_load=False): ...@@ -6,6 +6,10 @@ def load_torch_file(ckpt, safe_load=False):
import safetensors.torch import safetensors.torch
sd = safetensors.torch.load_file(ckpt, device="cpu") sd = safetensors.torch.load_file(ckpt, device="cpu")
else: else:
if safe_load:
if not 'weights_only' in torch.load.__code__.co_varnames:
print("Warning torch.load doesn't support weights_only on this pytorch version, loading unsafely.")
safe_load = False
if safe_load: if safe_load:
pl_sd = torch.load(ckpt, map_location="cpu", weights_only=True) pl_sd = torch.load(ckpt, map_location="cpu", weights_only=True)
else: else:
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment