Update on "Prevent legacy_load when weights_only=True (correctly)"

mikaylagawarecki · mikaylagawarecki · commit 9ec6b441c3ad · 2025-01-17T08:06:54.000-08:00
Only prevent `legacy_load` (.tar format removed in #713), not the whole of `_legacy_load` (.tar format + _use_new_zipfile_serialization=False) Differential Revision: [D68301405](https://our.internmc.facebook.com/intern/diff/D68301405) [ghstack-poisoned]
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -1612,6 +1612,11 @@ def persistent_load(saved_id):
         with closing(
             tarfile.open(fileobj=f, mode="r:", format=tarfile.PAX_FORMAT)
         ) as tar, mkdtemp() as tmpdir:
+            if pickle_module is _weights_only_unpickler:
+                raise RuntimeError(
+                    "Cannot use ``weights_only=True`` with files saved in the "
+                    "legacy .tar format. " + UNSAFE_MESSAGE
+                )
             tar.extract("storages", path=tmpdir)
             with open(os.path.join(tmpdir, "storages"), "rb", 0) as f:
                 num_storages = pickle_module.load(f, **pickle_load_args)
@@ -1738,13 +1743,6 @@ def persistent_load(saved_id):
         # legacy_load requires that f has fileno()
         # only if offset is zero we can attempt the legacy tar file loader
         try:
-            with closing(tarfile.open(fileobj=f, mode="r:", format=tarfile.PAX_FORMAT)):
-                if pickle_module is _weights_only_unpickler:
-                    raise RuntimeError(
-                        "Cannot use ``weights_only=True`` with files saved in the "
-                        "legacy .tar format. " + UNSAFE_MESSAGE
-                    )
-            f.seek(0)
             return legacy_load(f)
         except tarfile.TarError:
             if _is_zipfile(f):
