diff --git a/tensorflow_privacy/privacy/analysis/BUILD b/tensorflow_privacy/privacy/analysis/BUILD
index fc0812f3..08a9aac5 100644
--- a/tensorflow_privacy/privacy/analysis/BUILD
+++ b/tensorflow_privacy/privacy/analysis/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/dp_query/BUILD b/tensorflow_privacy/privacy/dp_query/BUILD
index 8c716e28..4d1e6804 100644
--- a/tensorflow_privacy/privacy/dp_query/BUILD
+++ b/tensorflow_privacy/privacy/dp_query/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/estimators/BUILD b/tensorflow_privacy/privacy/estimators/BUILD
index 851410ed..bcb21c12 100644
--- a/tensorflow_privacy/privacy/estimators/BUILD
+++ b/tensorflow_privacy/privacy/estimators/BUILD
@@ -1,4 +1,5 @@
 load("@rules_python//python:defs.bzl", "py_library")
+load("@rules_python//python:py_test.bzl", "py_test")
 
 package(
     default_visibility = ["//visibility:public"],
diff --git a/tensorflow_privacy/privacy/estimators/v1/BUILD b/tensorflow_privacy/privacy/estimators/v1/BUILD
index 463a015d..c6a7eb5b 100644
--- a/tensorflow_privacy/privacy/estimators/v1/BUILD
+++ b/tensorflow_privacy/privacy/estimators/v1/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/BUILD b/tensorflow_privacy/privacy/fast_gradient_clipping/BUILD
index 9dff250f..dc3b15f9 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/BUILD
+++ b/tensorflow_privacy/privacy/fast_gradient_clipping/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_applicable_licenses = ["//third_party/py/tensorflow_privacy:license"],
     default_visibility = ["//visibility:public"],
diff --git a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD
index 304a83dc..de331a12 100644
--- a/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD
+++ b/tensorflow_privacy/privacy/fast_gradient_clipping/registry_functions/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/keras_models/BUILD b/tensorflow_privacy/privacy/keras_models/BUILD
index 608d7e1b..709f6d58 100644
--- a/tensorflow_privacy/privacy/keras_models/BUILD
+++ b/tensorflow_privacy/privacy/keras_models/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_applicable_licenses = ["//third_party/py/tensorflow_privacy:license"],
     default_visibility = ["//visibility:public"],
diff --git a/tensorflow_privacy/privacy/logistic_regression/BUILD b/tensorflow_privacy/privacy/logistic_regression/BUILD
index 7835bc69..3161de3d 100644
--- a/tensorflow_privacy/privacy/logistic_regression/BUILD
+++ b/tensorflow_privacy/privacy/logistic_regression/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/optimizers/BUILD b/tensorflow_privacy/privacy/optimizers/BUILD
index 804f2af0..0d24d92b 100644
--- a/tensorflow_privacy/privacy/optimizers/BUILD
+++ b/tensorflow_privacy/privacy/optimizers/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/privacy_tests/BUILD b/tensorflow_privacy/privacy/privacy_tests/BUILD
index 19e877e2..bd4bd730 100644
--- a/tensorflow_privacy/privacy/privacy_tests/BUILD
+++ b/tensorflow_privacy/privacy/privacy_tests/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/BUILD b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/BUILD
index d5621c34..de1e14b4 100644
--- a/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/BUILD
+++ b/tensorflow_privacy/privacy/privacy_tests/membership_inference_attack/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/privacy_tests/secret_sharer/BUILD b/tensorflow_privacy/privacy/privacy_tests/secret_sharer/BUILD
index dbe141ea..9f55ae84 100644
--- a/tensorflow_privacy/privacy/privacy_tests/secret_sharer/BUILD
+++ b/tensorflow_privacy/privacy/privacy_tests/secret_sharer/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/sparsity_preserving_noise/BUILD b/tensorflow_privacy/privacy/sparsity_preserving_noise/BUILD
index c5c1befd..c2610002 100644
--- a/tensorflow_privacy/privacy/sparsity_preserving_noise/BUILD
+++ b/tensorflow_privacy/privacy/sparsity_preserving_noise/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )
diff --git a/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/BUILD b/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/BUILD
index 73e14813..f6d04d7c 100644
--- a/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/BUILD
+++ b/tensorflow_privacy/privacy/sparsity_preserving_noise/registry_functions/BUILD
@@ -1,3 +1,5 @@
+load("@rules_python//python:py_test.bzl", "py_test")
+
 package(
     default_visibility = ["//visibility:public"],
 )