diff --git a/tensorflow_gnn/runner/BUILD b/tensorflow_gnn/runner/BUILD
index 7e44459a..c0a4643f 100644
--- a/tensorflow_gnn/runner/BUILD
+++ b/tensorflow_gnn/runner/BUILD
@@ -1,9 +1,10 @@
+load("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_test")
+
 # Copybara rewrites load() statements back and forth; do not reformat.
 # buildifier: disable=out-of-order-load, disable=same-origin-load
 load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "pytype_strict_library")
 
 # buildifier: disable=out-of-order-load, disable=same-origin-load
-load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "py_strict_test")
 
 # buildifier: disable=out-of-order-load, disable=same-origin-load
 load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "distribute_py_test")
@@ -85,7 +86,7 @@ distribute_py_test(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "orchestration_test",
     srcs = ["orchestration_test.py"],
     deps = [
diff --git a/tensorflow_gnn/runner/input/BUILD b/tensorflow_gnn/runner/input/BUILD
index 8d6e73e8..045b249e 100644
--- a/tensorflow_gnn/runner/input/BUILD
+++ b/tensorflow_gnn/runner/input/BUILD
@@ -1,9 +1,10 @@
+load("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_test")
+
 # Copybara rewrites load() statements back and forth; do not reformat.
 # buildifier: disable=out-of-order-load, disable=same-origin-load
 load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "pytype_strict_library")
 
 # buildifier: disable=out-of-order-load, disable=same-origin-load
-load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "py_strict_test")
 
 licenses(["notice"])
 
@@ -22,7 +23,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "datasets_test",
     srcs = ["datasets_test.py"],
     deps = [
diff --git a/tensorflow_gnn/runner/tasks/BUILD b/tensorflow_gnn/runner/tasks/BUILD
index 0d35c5a0..c61a5f28 100644
--- a/tensorflow_gnn/runner/tasks/BUILD
+++ b/tensorflow_gnn/runner/tasks/BUILD
@@ -1,9 +1,10 @@
+load("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_test")
+
 # Copybara rewrites load() statements back and forth; do not reformat.
 # buildifier: disable=out-of-order-load, disable=same-origin-load
 load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "pytype_strict_library")
 
 # buildifier: disable=out-of-order-load, disable=same-origin-load
-load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "py_strict_test")
 
 licenses(["notice"])
 
@@ -23,7 +24,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "classification_test",
     srcs = ["classification_test.py"],
     deps = [
@@ -47,7 +48,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "link_prediction_test",
     srcs = ["link_prediction_test.py"],
     deps = [
@@ -70,7 +71,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "regression_test",
     srcs = ["regression_test.py"],
     deps = [
diff --git a/tensorflow_gnn/runner/utils/BUILD b/tensorflow_gnn/runner/utils/BUILD
index 5f3fc4b3..2c3d1d7d 100644
--- a/tensorflow_gnn/runner/utils/BUILD
+++ b/tensorflow_gnn/runner/utils/BUILD
@@ -1,3 +1,5 @@
+load("//third_party/bazel_rules/rules_python/python:py_test.bzl", "py_test")
+
 # Copybara rewrites load() statements back and forth; do not reformat.
 # buildifier: disable=out-of-order-load, disable=same-origin-load
 load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "pytype_strict_binary")
@@ -6,7 +8,6 @@ load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "pytype_strict_binary
 load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "pytype_strict_library")
 
 # buildifier: disable=out-of-order-load, disable=same-origin-load
-load("@tensorflow_gnn//tensorflow_gnn:tensorflow_gnn.bzl", "py_strict_test")
 
 licenses(["notice"])
 
@@ -95,7 +96,7 @@ pytype_strict_library(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "attribution_test",
     srcs = ["attribution_test.py"],
     deps = [
@@ -107,7 +108,7 @@ py_strict_test(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "model_export_test",
     srcs = ["model_export_test.py"],
     deps = [
@@ -120,7 +121,7 @@ py_strict_test(
     ],
 )
 
-py_strict_test(
+py_test(
     name = "parsing_test",
     srcs = ["parsing_test.py"],
     deps = [