ClearML/clearml
1
0
Fork 0
mirror of https://github.com/clearml/clearml synced 2025-03-05 03:29:25 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix Pytorch ScriptModule autobind

Browse Source
This commit is contained in:
allegroai 2022-04-13 14:15:00 +03:00
parent 681d75a309
commit 6695c94fdd
1 changed files with 6 additions and 0 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

6
clearml/binding/frameworks/pytorch_bind.py
View File

@ -38,6 +38,12 @@ class PatchPyTorchModelIO(PatchBaseModelIO):
import torch # noqa
torch.save = _patched_call(torch.save, PatchPyTorchModelIO._save)
torch.load = _patched_call(torch.load, PatchPyTorchModelIO._load)
# noinspection PyBroadException
try:
# noinspection PyProtectedMember
torch.jit._script.RecursiveScriptModule.save = _patched_call(torch.jit._script.RecursiveScriptModule.save, PatchPyTorchModelIO._save)
except BaseException:
pass
# no need to worry about recursive calls, _patched_call takes care of that
if hasattr(torch, 'serialization') and hasattr(torch.serialization, '_save'):