diff --git a/src/nnsight/envoy.py b/src/nnsight/envoy.py index b1f2755e..d498a59e 100644 --- a/src/nnsight/envoy.py +++ b/src/nnsight/envoy.py @@ -235,7 +235,7 @@ def modules( if included: if names: - envoys.append((self._module_path, self)) + envoys.append((self._module_path.lstrip("."), self)) else: envoys.append(self) diff --git a/src/nnsight/models/NNsightModel.py b/src/nnsight/models/NNsightModel.py index 32d32204..2f39479c 100644 --- a/src/nnsight/models/NNsightModel.py +++ b/src/nnsight/models/NNsightModel.py @@ -300,6 +300,9 @@ def __setattr__(self, key: Any, value: Any) -> None: if key not in ('_model', '_model_key') and isinstance(value, torch.nn.Module): + if self._envoy._tracer is None: + setattr(self._model, key, value) + setattr(self._envoy, key, value) else: @@ -312,6 +315,9 @@ def __getattr__(self, key: Any) -> Union[Envoy, InterventionProxy, Any]: Returns: Any: Attribute. """ + if self._envoy._tracer is None: + return getattr(self._model, key) + return getattr(self._envoy, key) ### NNsight VIRTUAL METHODS BELOW #####################################