Reflect review comments

tensorflow · Oct 28, 2024 · b09ac2a · b09ac2a
1 parent 33724c1
commit b09ac2a
Show file tree

Hide file tree

Showing 3 changed files with 12 additions and 3 deletions.
diff --git a/.github/workflows/ci-test.yml b/.github/workflows/ci-test.yml
@@ -62,4 +62,4 @@ jobs:
     - name: Run unit tests
       shell: bash
       run: |
-        pytest -vv -m "${{ matrix.which-tests }}" tfx
+        pytest -m "${{ matrix.which-tests }}"
diff --git a/...al/distributed_inference/graphdef_experiments/subgraph_partitioning/beam_pipeline_test.py b/...al/distributed_inference/graphdef_experiments/subgraph_partitioning/beam_pipeline_test.py
@@ -25,7 +25,11 @@
 from tfx.experimental.distributed_inference.graphdef_experiments.subgraph_partitioning import create_complex_graph
 from tfx.experimental.distributed_inference.graphdef_experiments.subgraph_partitioning import graph_partition
 
-tf.compat.v1.enable_eager_execution()  # Re-enable eager mode
+@pytest.fixture
+def enable_eager_execution_at_test_end():
+  yield
+  # Re-enable eager mode that create_complex_graph diabled.
+  tf.compat.v1.enable_eager_execution()
 
 class BeamPipelineTest(tf.test.TestCase):
   """A test for the beam pipeline library."""

diff --git a/.../distributed_inference/graphdef_experiments/subgraph_partitioning/graph_partition_test.py b/.../distributed_inference/graphdef_experiments/subgraph_partitioning/graph_partition_test.py
@@ -23,7 +23,12 @@
 from tfx.experimental.distributed_inference.graphdef_experiments.subgraph_partitioning import graph_partition
 from google.protobuf import text_format
 
-tf.compat.v1.enable_eager_execution()  # Re-enable eager mode
+@pytest.fixture
+def enable_eager_execution_at_test_end():
+  yield
+  # Re-enable eager mode that create_complex_graph diabled.
+  tf.compat.v1.enable_eager_execution()
+
 
 class RemoteOpLayerTest(tf.test.TestCase):
   """A test for the class _RemoteOpLayer."""