pytorch/pytorch

Jitted function cannot be pickled

Chiusa

#61.210 aperta il 3 lug 2021

 (10 commenti) (0 reazioni) (1 assegnatario)Python (28.880 fork)batch import
OSS contribution wantedgood first issueoncall: jit

Metriche repository

Star
 (102.396 stelle)
Metriche merge PR
 (Merge medio 1g 10h) (42 PR mergiate in 30 g)

Descrizione

🚀 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

Guida contributor