def METHOD_NAME(ctx_factory, func):
    """
    Wrap each generator invocation with the context manager factory.
    The input should be a function that returns a context manager,
    not a context manager itself, to handle one-shot context managers.
    """
    @functools.wraps(func)
    def generator_context(*args, **kwargs):
        gen = func(*args, **kwargs)
        # Generators are suspended and unsuspended at `yield`, hence we
        # make sure the grad mode is properly set every time the execution
        # flow returns into the wrapped generator and restored when it
        # returns through our `yield` to our caller (see PR #49017).
        try:
            # Issuing `None` to a generator fires it up
            with ctx_factory():
                response = gen.send(None)
            while True:
                try:
                    # Forward the response to our caller and get its next request
                    request = yield response
                except GeneratorExit:
                    # Inform the still active generator about its imminent closure
                    with ctx_factory():
                        gen.close()
                    raise
                except BaseException:
                    # Propagate the exception thrown at us by the caller
                    with ctx_factory():
                        response = gen.throw(*sys.exc_info())
                else:
                    # Pass the last request to the generator and get its response
                    with ctx_factory():
                        response = gen.send(request)
        # We let the exceptions raised above by the generator's `.throw` or
        # `.send` methods bubble up to our caller, except for StopIteration
        except StopIteration as e:
            # The generator informed us that it is done: take whatever its
            # returned value (if any) was and indicate that we're done too
            # by returning it (see docs for python's return-statement).
            return e.value
    return generator_context