move routing info to avoid circ import

microsoft · Aug 8, 2024 · 77a4ab7 · 77a4ab7
1 parent 8c9a02c
commit 77a4ab7
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/mttl/models/expert_context.py b/mttl/models/expert_context.py
@@ -2,8 +2,6 @@
 import threading
 from typing import List
 
-from mttl.models.modifiers.routing import RoutingInfo
-
 
 class InfoContainer:
     local = threading.local()
@@ -36,7 +34,7 @@ def routing_gates(self):
         return self._routing_gates
 
     @routing_infos.setter
-    def routing_infos(self, value: RoutingInfo):
+    def routing_infos(self, value: "RoutingInfo"):
         self._routing_infos = value
 
     @routing_gates.setter
@@ -48,6 +46,7 @@ def wrap_forward(cls, f):
         """
         Decorator method that wraps a ``forward()`` function of a model class.
         """
+        from mttl.models.modifiers.routing import RoutingInfo
 
         @functools.wraps(f)
         def wrapper_func(model, *args, **kwargs):