diff --git a/tensorflow_gnn/models/mt_albis/BUILD b/tensorflow_gnn/models/mt_albis/BUILD
index 5accc114..1bb15f7e 100644
--- a/tensorflow_gnn/models/mt_albis/BUILD
+++ b/tensorflow_gnn/models/mt_albis/BUILD
@@ -7,16 +7,19 @@ package(
     default_visibility = [":__subpackages__"],
 )
 
-package_group(name = "users")
+package_group(
+    name = "users",
+    includes = [
+        "public",
+    ],
+)
 
+# This is the public BUILD target for users of this model.
 pytype_strict_library(
     name = "mt_albis",
     srcs = ["__init__.py"],
     srcs_version = "PY3",
-    visibility = [
-        ":__subpackages__",
-        ":users",
-    ],
+    visibility = ["//visibility:public"],
     deps = [
         ":config_dict",
         ":layers",
@@ -29,10 +32,7 @@ pytype_strict_library(
 pytype_strict_library(
     name = "hparams_vizier",
     srcs = ["hparams_vizier.py"],
-    visibility = [
-        ":__subpackages__",
-        ":users",
-    ],
+    visibility = ["//visibility:public"],
     deps = [
         "//:expect_vizier_service_pyvizier_installed",
     ],
@@ -40,9 +40,6 @@ pytype_strict_library(
 
 exports_files(
     srcs = ["hparams_vizier.py"],
-    visibility = [
-        ":__subpackages__",
-    ],
 )
 
 # =============================================================================