From d2e4062607aef51e9b82be2c9989f3f19e5d0830 Mon Sep 17 00:00:00 2001 From: shavit Date: Sat, 31 Aug 2024 06:15:24 -0400 Subject: [PATCH] fix(io): load weights only in torch.load --- trainer/io.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/trainer/io.py b/trainer/io.py index 165f6ea..615ad0a 100644 --- a/trainer/io.py +++ b/trainer/io.py @@ -77,10 +77,10 @@ def load_fsspec( filecache={"cache_storage": str(get_user_data_dir("tts_cache"))}, mode="rb", ) as f: - return torch.load(f, map_location=map_location, **kwargs) + return torch.load(f, map_location=map_location, weights_only=True, **kwargs) else: with fsspec.open(str(path), "rb") as f: - return torch.load(f, map_location=map_location, **kwargs) + return torch.load(f, map_location=map_location, weights_only=True, **kwargs) def load_checkpoint(