diff --git a/.github/workflows/ci-test.yml b/.github/workflows/ci-test.yml
index a2982c92c6..3e3d9b95c0 100644
--- 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/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/beam_pipeline_test.py b/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/beam_pipeline_test.py
index ed64e35ee2..379beea02c 100644
--- a/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/beam_pipeline_test.py
+++ b/tfx/experimental/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/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/graph_partition_test.py b/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/graph_partition_test.py
index 56a04dac62..ca5e34d8a6 100644
--- a/tfx/experimental/distributed_inference/graphdef_experiments/subgraph_partitioning/graph_partition_test.py
+++ b/tfx/experimental/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."""