TorchModel now has graph description in meta

Oxid15 · Jul 29, 2022 · 580208a · 580208a
1 parent 375894a
commit 580208a
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