Use TF2 for mwms_pjrt_gpu_test

PiperOrigin-RevId: 635898879
tensorflow · May 21, 2024 · 6af67b7 · 6af67b7
1 parent 68dd380
commit 6af67b7
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/tensorflow/python/distribute/BUILD b/tensorflow/python/distribute/BUILD
@@ -658,13 +658,13 @@ cuda_py_strict_test(
     srcs_version = "PY3",
     tags = [
         "multi_and_single_gpu",
-        "notap",  # TODO(b/341375925): Re-enable this test when flakiness is fixed.
     ],
     xla_enabled = True,
     deps = [
         ":multi_process_runner",
         ":multi_worker_test_base",
         "//tensorflow/core:protos_all_py",
+        "//tensorflow/python/compat:v2_compat",
         "//tensorflow/python/distribute/cluster_resolver:cluster_resolver_lib",
         "//tensorflow/python/eager:context",
         "//tensorflow/python/eager:def_function",

diff --git a/tensorflow/python/distribute/mwms_pjrt_gpu_test.py b/tensorflow/python/distribute/mwms_pjrt_gpu_test.py
@@ -16,6 +16,7 @@
 import copy
 
 from tensorflow.core.protobuf import tensorflow_server_pb2
+from tensorflow.python.compat import v2_compat
 from tensorflow.python.distribute import cluster_resolver as cluster_resolver_lib
 from tensorflow.python.distribute import multi_process_runner
 from tensorflow.python.distribute import multi_worker_test_base
@@ -113,4 +114,5 @@ def f(x):
 
 
 if __name__ == "__main__":
+  v2_compat.enable_v2_behavior()
   multi_process_runner.test_main()