def METHOD_NAME(*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')