Fix from last update...

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)