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
@@ -377,15 +377,15 @@ class VQAutoEncoder(nn.Module):
)
if model_path is not None:
chkpt = torch.load(model_path, map_location="cpu")
chkpt = torch.load(model_path, map_location="cpu", weights_only=True)
if "params_ema" in chkpt:
self.load_state_dict(
torch.load(model_path, map_location="cpu")["params_ema"]
torch.load(model_path, map_location="cpu", weights_only=True)["params_ema"]
)
logger.info(f"vqgan is loaded from: {model_path} [params_ema]")
elif "params" in chkpt:
self.load_state_dict(
torch.load(model_path, map_location="cpu")["params"]
torch.load(model_path, map_location="cpu", weights_only=True)["params"]
)
logger.info(f"vqgan is loaded from: {model_path} [params]")
else: