From fc7e8e2854ccf808156b13517ab592674026d75b Mon Sep 17 00:00:00 2001 From: Constantin Date: Mon, 30 Dec 2024 23:31:29 +0900 Subject: [PATCH] Compatiblity with PyTorch 2.4+ (weights_only default value changed) Starting from version 2.4 PyTorch introduces a stricter check for the objects which can be loaded with torch.load(); to use bark successfully with PyTorch>=2.4 the weights_only attribute needs to be set explicitly --- bark/generation.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/bark/generation.py b/bark/generation.py index 54f98709..5b0906e4 100644 --- a/bark/generation.py +++ b/bark/generation.py @@ -209,7 +209,7 @@ def _load_model(ckpt_path, device, use_small=False, model_type="text"): if not os.path.exists(ckpt_path): logger.info(f"{model_type} model not found, downloading into `{CACHE_DIR}`.") _download(model_info["repo_id"], model_info["file_name"]) - checkpoint = torch.load(ckpt_path, map_location=device) + checkpoint = torch.load(ckpt_path, map_location=device, weights_only=False) # this is a hack model_args = checkpoint["model_args"] if "input_vocab_size" not in model_args: