From 6695c94fdd8b404be23a7b6f607ccdcc2b2e4b14 Mon Sep 17 00:00:00 2001
From: allegroai <>
Date: Wed, 13 Apr 2022 14:15:00 +0300
Subject: [PATCH] Fix Pytorch ScriptModule autobind

---
 clearml/binding/frameworks/pytorch_bind.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/clearml/binding/frameworks/pytorch_bind.py b/clearml/binding/frameworks/pytorch_bind.py
index 5e4007cb..1a7c92e4 100644
--- a/clearml/binding/frameworks/pytorch_bind.py
+++ b/clearml/binding/frameworks/pytorch_bind.py
@@ -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'):