Added ScriptFunction pkl exception for issue #61210 #61381 (#67076)

Summary:
Fixes https://github.com/pytorch/pytorch/issues/61381, https://github.com/pytorch/pytorch/issues/61210

Pull Request resolved: https://github.com/pytorch/pytorch/pull/67076

Reviewed By: jbschlosser

Differential Revision: D32908175

Pulled By: suo

fbshipit-source-id: f6e175793243dc96cde5e44022d92f2623b934eb

Co-authored-by: LucaStubbe <stubbeluca@gmail.com>
Co-authored-by: Kanon Tromp <ktromp1@student.cccd.edu>
This commit is contained in:
Nik B
2021-12-09 09:42:55 -08:00
committed by Facebook GitHub Bot
parent d21646c432
commit 2d5b3101c1
2 changed files with 16 additions and 0 deletions

View File

@ -56,6 +56,13 @@ function and does not have any attributes or Parameters.
"""
set_module(ScriptFunction, "torch.jit")
# Throws an error if a jit function is pickled.
# Helps to avoid Python crashes for Python versions 3.9.5 + when protocol 0 or 1 is given as an argument.
def _reduce(cls):
raise pickle.PickleError("ScriptFunction cannot be pickled")
ScriptFunction.__reduce__ = _reduce # type: ignore[assignment]
if _enabled:
Attribute = collections.namedtuple("Attribute", ["value", "type"])