Add example for torch.serialization.add_safe_globals

mikaylagawarecki · mikaylagawarecki · commit cfcc420a1ca3 · 2024-06-25T07:58:46.000-07:00
ghstack-source-id: 6dc3275 Pull Request resolved: #129396 (cherry picked from commit ed8c36eda0f4dcf7b1d9c5eb2fb1cdccdf3fee6e)
diff --git a/torch/serialization.py b/torch/serialization.py
@@ -172,7 +172,23 @@ def add_safe_globals(safe_globals: List[Any]) -> None:
 
     Args:
         safe_globals (List[Any]): list of globals to mark as safe
-    '''
+
+    Example:
+        >>> # xdoctest: +SKIP("Can't torch.save(t, ...) as doctest thinks MyTensor is defined on torch.serialization")
+        >>> import tempfile
+        >>> class MyTensor(torch.Tensor):
+        ...     pass
+        >>> t = MyTensor(torch.randn(2, 3))
+        >>> with tempfile.NamedTemporaryFile() as f:
+        ...     torch.save(t, f.name)
+        # Running `torch.load(f.name, weights_only=True)` will fail with
+        # Unsupported global: GLOBAL __main__.MyTensor was not an allowed global by default.
+        # Check the code and make sure MyTensor is safe to be used when loaded from an arbitrary checkpoint.
+        ...     torch.serialization.add_safe_globals([MyTensor])
+        ...     torch.load(f.name, weights_only=True)
+        # MyTensor([[-0.5024, -1.8152, -0.5455],
+        #          [-0.8234,  2.0500, -0.3657]])
+    """
     _weights_only_unpickler._add_safe_globals(safe_globals)
 
 def _is_zipfile(f) -> bool:
