Pass forward pass args to context

src/transformers/adapters/context.py

@@ -72,10 +72,10 @@ class ForwardContext:
  # thread-local storage that holds a stack of active contexts
  storage = threading.local()
 
- def __init__(self, model):
+ def __init__(self, model, *args, **kwargs):
  # If the model has a method ``forward_context()``, use it to create the context.
  if hasattr(model, "forward_context"):
- model.forward_context(self)
+ model.forward_context(self, *args, **kwargs)
 
  @classmethod
  def wrap(cls, f):
@@ -85,7 +85,7 @@ def wrap(cls, f):
 
  @functools.wraps(f)
  def wrapper_func(self, *args, **kwargs):
- context = cls(self)
+ context = cls(self, *args, **kwargs)
  cls.get_contexts().append(context)
  results = f(self, *args, **kwargs)
  cls.get_contexts().pop()

src/transformers/adapters/model_mixin.py

@@ -563,7 +563,7 @@ def freeze_model(self, freeze=True):
  param.requires_grad = not freeze
  self.model_freezed = freeze
 
- def forward_context(self, context: ForwardContext):
+ def forward_context(self, context: ForwardContext, *args, **kwargs):
  """
  This method is called by the ``ForwardContext`` at the beginning of the forward pass.
  """