From bec660ba2d83ef5a05aa690083983a93ae967984 Mon Sep 17 00:00:00 2001
From: jcgu <jcgu@google.com>
Date: Mon, 24 Aug 2026 14:17:20 -0700
Subject: [PATCH] Fix PyTorch missing visibility compilation failures

PiperOrigin-RevId: 970075544
---
 tpu_sync/frameworks/torch/BUILD | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tpu_sync/frameworks/torch/BUILD b/tpu_sync/frameworks/torch/BUILD
index 5f21dab4..618bd2b0 100644
--- a/tpu_sync/frameworks/torch/BUILD
+++ b/tpu_sync/frameworks/torch/BUILD
@@ -429,11 +429,11 @@ cc_library(
     tags = ["nobuilder"],
     visibility = ["//visibility:public"],
     deps = [
-        ":torch_tpu_tensor_to_buffer_headers",
         ":torch_utils_mock",
         "//tpu_sync/weight_sync:weight_synchronizer_base",
         "@torch_tpu//shims/torch:aten_headers",
         "@torch_tpu//shims/torch:torch_headers",
+        "@torch_tpu//torch_tpu/eager:device_buffer",
     ],
 )