Merge pull request #75 from Oxid15/torch_model_update

TorchModel now has graph description in meta
Oxid15 · Jul 30, 2022 · 2372418 · 2372418
2 parents 8bd1dc6 + 580208a
commit 2372418
Showing 1 changed file with 5 additions and 0 deletions.
diff --git 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