def _v2_distributed_test(f, test_or_class, *args, **kwargs): with context.eager_mode(): with testing_utils.run_eagerly_scope(False): with testing_utils.run_distributed_scope(True): f(test_or_class, *args, **kwargs)
def _v1_graph_test(f, test_or_class, config, *args, **kwargs): with context.graph_mode(), testing_utils.run_eagerly_scope(False): with testing_utils.run_distributed_scope(False): with test_or_class.test_session(use_gpu=True, config=config): f(test_or_class, *args, **kwargs)