From 580208ab50ddf3f7607c08a254826d33b67eded6 Mon Sep 17 00:00:00 2001
From: Ilia Moiseev <villeman.5@yandex.ru>
Date: Fri, 29 Jul 2022 18:56:41 +0300
Subject: [PATCH] TorchModel now has graph description in meta

---
 cascade/utils/torch_model.py | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/cascade/utils/torch_model.py b/cascade/utils/torch_model.py
index bb3a9507..026c7916 100644
--- a/cascade/utils/torch_model.py
+++ b/cascade/utils/torch_model.py
@@ -34,3 +34,8 @@ def save(self, path, *args, **kwargs) -> None:
     def load(self, path, *args, **kwargs) -> None:
         with open(path, 'rb') as f:
             self._model = torch.load(f)
+
+    def get_meta(self):
+        meta = super().get_meta()
+        meta[-1]['module'] = repr(self._model)
+        return meta