test_utils.py 文件源码

def keras_test(func):
    """Function wrapper to clean up after TensorFlow tests.

    # Arguments
        func: test function to clean up after.

    # Returns
        A function wrapping the input function.
    """
    @six.wraps(func)
    def wrapper(*args, **kwargs):
        output = func(*args, **kwargs)
        if K.backend() == 'tensorflow':
            K.clear_session()
        return output
    return wrapper