From 77a4ab7ce9523ab73b83f169274e1155f50d1b07 Mon Sep 17 00:00:00 2001 From: Alessandro Sordoni Date: Thu, 8 Aug 2024 12:52:24 -0700 Subject: [PATCH] move routing info to avoid circ import --- mttl/models/expert_context.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/mttl/models/expert_context.py b/mttl/models/expert_context.py index 4669231b3..c44f2daeb 100644 --- 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):