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