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",
     ],
 )