consider wrappers such as nn.DataParallel. Ignoring the other specifics of what that class is doing, the basic pattern is that you have a class which simply wraps the function of a base module:
class Wrapper(nn.Module): def __init__(self, base_module: nn.Module): super().__init__() self.base_module = base_module
the real trouble comes with generically wrapping a forward pass:
def forward(self, *args, **kwargs): self.base_module(*args, **kwargs)
I would love to see this wrapper itself be jit scriptable so long as the base_module is. In my particular use case, I actually don’t need to add anything to this forward pass. I know that I may simply have to require the forward method to only take a single argument, but I would like to see if anyone has handled this in other ways.