Merge pull request #263 from ndif-team/calling

Add ability to call envoys /  models outside of tracing context

src/nnsight/envoy.py

@@ -142,12 +142,28 @@ def _set_tracer(self, tracer: Tracer, propagate=True):
         if propagate:
             for envoy in self._sub_envoys:
                 envoy._set_tracer(tracer, propagate=True)
+
+
+    def _tracing(self) -> bool:
+        """Whether or not tracing.
+
+        Returns:
+            bool: Is tracing.
+        """
+
+        try:
+
+            return self._tracer.graph.alive
+
+        except:
+
+            return False
 
     def _scanning(self) -> bool:
         """Whether or not in scanning mode. Checks the current Tracer's Invoker.
 
         Returns:
-            bool: _description_
+            bool: Is scanning.
         """
 
         try:
@@ -420,6 +436,9 @@ def __call__(
         Returns:
             InterventionProxy: Module call proxy.
         """
+
+        if not self._tracing():
+            return self._module(*args, **kwargs)
 
         if isinstance(self._tracer.backend, EditBackend):
             hook = True

src/nnsight/models/NNsightModel.py

@@ -119,6 +119,10 @@ def __init__(
             self.dispatch_model()
 
         logger.info(f"Initialized `{self._model_key}`")
+
+    def __call__(self, *args, **kwargs):
+
+        return self._envoy(*args, **kwargs)
 
     def trace(
         self,