pytorch
diff --git a/‎torch/_weights_only_unpickler.py
Lines changed: 26 additions & 3 deletions b/‎torch/_weights_only_unpickler.py
Lines changed: 26 additions & 3 deletions
diff --git a/‎torch/serialization.py
Lines changed: 3 additions & 1 deletion b/‎torch/serialization.py
Lines changed: 3 additions & 1 deletion
@@ -23,7 +23,7 @@
 # weights = torch.load(buf, weights_only = True)
 
 import functools as _functools
-from _compat_pickle import IMPORT_MAPPING, NAME_MAPPING
+import warnings
 from collections import Counter, OrderedDict
 from pickle import (
     APPEND,
@@ -68,6 +68,16 @@
 from sys import maxsize
 from typing import Any, Dict, List
 
+try:
+    # We rely on this module in private cPython which provides dicts of
+    # modules/functions that had their names changed from Python 2 to 3
+    has_compat_pickle = True
+    from _compat_pickle import IMPORT_MAPPING, NAME_MAPPING
+except ImportError:
+    # To prevent warning on import torch, we warn in the Unpickler.load below
+    has_compat_pickle = False
+    IMPORT_MAPPING, NAME_MAPPING = dict(), dict()
+
 import torch
 
 _marked_safe_globals_list: List[Any] = []
@@ -179,6 +189,13 @@ def load(self):
 
         Return the reconstituted object hierarchy specified in the file.
         """
+        if not has_compat_pickle:
             warnings.warn(
+                "Could not import IMPORT_MAPPING and NAME_MAPPING from _compat_pickle. "
+                "If the default `pickle_protocol` was used at `torch.save` time, any functions or "
+                "classes that are in these maps might not behave correctly if allowlisted via "
+                "`torch.serialization.add_safe_globals()`."
+            )
         self.metastack = []
         self.stack: List[Any] = []
         self.append = self.stack.append
@@ -195,7 +212,7 @@ def load(self):
                 name = readline()[:-1].decode("utf-8")
                 # Patch since torch.save default protocol is 2
                 # users will be running this code in python > 3
-                if self.proto == 2:
+                if self.proto == 2 and has_compat_pickle:
                     if (module, name) in NAME_MAPPING:
                         module, name = NAME_MAPPING[(module, name)]
                     elif module in IMPORT_MAPPING:
@@ -352,8 +369,14 @@ def load(self):
                 self.append(decode_long(data))
             # First and last deserializer ops
             elif key[0] == PROTO[0]:
-                # Read and ignore proto version
                 self.proto = read(1)[0]
+                if self.proto != 2:
+                    warnings.warn(
+                        f"Detected pickle protocol {self.proto} in the checkpoint, which was "
+                        "not the default pickle protocol used by `torch.load` (2). The weights_only "
+                        "Unpickler might not support all instructions implemented by this protocol, "
+                        "please file an issue for adding support if you encounter this."
+                    )
             elif key[0] == STOP[0]:
                 rc = self.stack.pop()
                 return rc
 
@@ -203,7 +203,9 @@ def get_safe_globals() -> List[Any]:
 
 def add_safe_globals(safe_globals: List[Any]) -> None:
     """
-    Marks the given globals as safe for ``weights_only`` load.
+    Marks the given globals as safe for ``weights_only`` load. For example, functions
+    added to this list can be called during unpickling, classes could be instantiated
+    and have state set.
 
     Args:
         safe_globals (List[Any]): list of globals to mark as safe