def METHOD_NAME(module_or_function: Union[nn.Module, Callable[..., Any]],
                 with_weights: bool = True):
    """
    Compile a Python function with PyTorch to an LBANN graph lazily. This means
    that whenever the decorated function is called, an LBANN graph will be
    returned instead of calling the function.
    :param module_or_function: ``torch.nn.Module`` or function to mark for
                               compilation.
    :param with_weights: If True (default), also stores the parameters of
                         the given model as constant initializers of the LBANN
                         graph's weights.
    """
    converters.load_replacements()
    argnames = _get_module_argnames(module_or_function)
    f = module_or_function
    if isinstance(f, nn.Module):
        f = opaque.wrap_opaque_modules(f)
        print('Replaced opaque operators on the module tree:',
              opaque.count_replaced_submodules(f))
    f = dynamo.optimize(functools.partial(lowering.dynamo_callback, argnames,
                                          with_weights),
                        nopython=True)(f)
    @functools.wraps(f)
    def _wrapped(*args, **kwargs):
        try:
            f(*args, **kwargs)
        except BackendCompilerFailed as ex:
            if isinstance(ex.inner_exception, LBANNGraph):
                return ex.inner_exception.graph
            raise
        except LBANNGraph as ex:
            return ex.graph
        raise RuntimeError(
            'Could not extract an LBANN graph from PyTorch module '
            'or function')
    return _wrapped