From 416a6f24925bbfa5b2efad8315522d904664a46d Mon Sep 17 00:00:00 2001 From: "jadenfk@outlook.com" Date: Tue, 28 May 2024 17:23:53 -0400 Subject: [PATCH] Fix from last update... --- src/nnsight/models/NNsightModel.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/src/nnsight/models/NNsightModel.py b/src/nnsight/models/NNsightModel.py index 56c370d1..32d32204 100644 --- a/src/nnsight/models/NNsightModel.py +++ b/src/nnsight/models/NNsightModel.py @@ -76,9 +76,7 @@ def __init__( if dispatch and not self._dispatched: # Dispatch ._model on initialization vs lazy dispatching. self.dispatch_model() - - self._initialized = True - + logger.info(f"Initialized `{self._model_key}`") def trace( @@ -300,7 +298,7 @@ def __repr__(self) -> str: def __setattr__(self, key: Any, value: Any) -> None: """Overload setattr to create and set an Envoy when trying to set a torch Module.""" - if '_initialized' in self.__dict__ and isinstance(value, torch.nn.Module): + if key not in ('_model', '_model_key') and isinstance(value, torch.nn.Module): setattr(self._envoy, key, value)