A restricted unpickler is a pickle reader that refuses to import anything not on an allow list. It does this by overriding the find_class method, so a model file asking for os.system or eval fails with an error instead of running. It narrows the risk a lot, but it is not a guarantee.
Why pickle needs restricting at all
Python's pickle format is a list of instructions. Some of them tell the reader to import a function and call it. The Python documentation says it plainly: the pickle module is not secure, and malicious pickle data can execute arbitrary code during unpickling. Many PyTorch checkpoints are pickles, which is why a downloaded model file can be a code execution path.
The Python docs describe the control point. By default, unpickling imports any class or function found in the data. You can change that by subclassing the unpickler and overriding find_class, which is called whenever a global (a class or function) is requested.
How the allow list works
Every time the file needs a global, the reader calls find_class with a module name and an object name. Your override decides what happens next.
import builtins
import io
import pickle
SAFE = {("builtins", "range"), ("builtins", "set")}
class RestrictedUnpickler(pickle.Unpickler):
def find_class(self, module, name):
if (module, name) in SAFE:
return getattr(builtins, name)
raise pickle.UnpicklingError(f"global {module}.{name} is forbidden")
def restricted_loads(data):
return RestrictedUnpickler(io.BytesIO(data)).load()
This is the same pattern the Python docs show. A payload that references os.system is rejected before anything runs. A file that only needs plain containers and numbers loads normally.
What this means for model files
A real model checkpoint needs more than builtins. It references tensor types and a few helper functions. PyTorch's own loader applies this idea: with weights_only set, it limits loading to tensors, primitive types and dictionaries, plus any types you register as safe. That list is short on purpose.
The practical lesson for a team:
| Approach | What it does | Main limit |
|---|---|---|
| Plain unpickle | Imports anything the file asks for | Runs attacker code |
| Restricted unpickler | Allows only listed globals | Allow list must be tight and kept current |
| Weights only loading | Built in allow list for tensor data | Only covers tensor style content |
| Non pickle format | No code path in the format | Needs the publisher to provide it |
Where it falls short
The Python docs add a caution after the example: you have to be careful with what you allow. Every entry on the list is code you are choosing to trust. Allow a function that can open files, and you have reopened the door.
Three more limits are worth knowing:
- An allow list guards imports. It does not tell you whether the tensor values were poisoned or the model behaves badly.
- Teams tend to widen the list when a load fails, and it drifts. Review changes to it like any security rule.
- It only protects code paths that actually use your unpickler. A notebook that calls plain pickle.load on the same file is still exposed.
If you control both ends, the cleaner answer is to avoid pickle for data that crosses a trust boundary. The Python docs suggest json for untrusted data, and the model world has formats such as safetensors that carry only tensor data and a header.
What to do this week
- Search your repos for pickle.load, pickle.loads and torch.load calls.
- Replace them with a restricted unpickler or a weights only load where the file source is not fully trusted.
- Prefer a non pickle weights file when the publisher offers one.
- Record the source and digest of every model file you load, so a bad file can be traced.
Quick check
Does find_class stop all malicious pickles?
No. It stops imports you did not allow. It cannot protect you from a permissive allow list, from code paths that bypass it, or from poisoned weight values.
Is torch.load a restricted unpickler?
With weights_only enabled, it applies a similar allow list limited to tensors, primitive types and dictionaries, plus types you register. Never load data from an untrusted source with it disabled.
Safeguard scans your repositories for dependencies and builds SBOMs, and policies and gates can block risky changes in pull requests. You can try a free scan at gold.safeguard.sh/scan.
Sources
Last reviewed: 19 October 2026.
