Function patch_forward

exir/capture/_capture.py:139–155

Helper method to make it easier to cleanly torch.export() a method on a module that is not `forward`. TODO(suo): upstream this to torch.export.wrapper.

(obj: torch.nn.Module, new_method)

@contextmanager
def patch_forward(obj: torch.nn.Module, new_method):
    """Helper method to make it easier to cleanly torch.export() a method on a
    module that is not `forward`.

    TODO(suo): upstream this to torch.export.wrapper.
    """
    # Save the original method
    original_method = obj.forward

    # Patch the method
    obj.forward = new_method.__get__(obj, obj.__class__)

    try:
        yield
    finally:
        # Restore the original method
        obj.forward = original_method


class WrapperModule(torch.nn.Module):