pytorch/pytorch

Jitted function cannot be pickled

Geschlossen

#61.210 geöffnet am 03.07.2021

 (10 Kommentare) (0 Reaktionen) (1 zugewiesene Person)Python (28.910 Forks)batch import
OSS contribution wantedgood first issueoncall: jit

Repository-Metriken

Stars
 (102.472 Sterne)
PR-Merge-Metriken
 (Durchschn. Merge 1T 10h) (42 gemergte PRs in 30 T)

Beschreibung

🚀 Feature

Support to pickle a jitted function (or at least throw a TypeError when using protocol 0 and 1).

Motivation

Trying to pickle a jitted function either raises TypeError: cannot pickle 'torch._C.ScriptFunction' object when protocol>1 or far worse when using protocol=0 or protocol=1 python 3.9.5 dies with:

terminate called after throwing an instance of 'std::runtime_error'
  what():  instance allocation failed: new instance has no pybind11-registered base types
Aborted (core dumped)

Example:

import pickle

import torch

@torch.jit.script
def fun(x: torch.Tensor) -> torch.Tensor:
    return x

pickle.dumps(fun, protocol=0)

Pitch

Implement pickle support for all pickle protocols (or raise a TypeError when used with protocl<2).

Additional context

I encountered this when using a jitted function with dask https://github.com/dask/dask/issues/7861

Contributor Guide