fix(io): load weights only in torch.load

idiap · Sep 12, 2024 · d2e4062 · d2e4062
1 parent 8025277
commit d2e4062
Showing 1 changed file with 2 additions and 2 deletions.
diff --git 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(