diff --git a/tpu_sync/frameworks/torch/BUILD b/tpu_sync/frameworks/torch/BUILD
index 5f21dab4..73afb299 100644
--- a/tpu_sync/frameworks/torch/BUILD
+++ b/tpu_sync/frameworks/torch/BUILD
@@ -527,8 +527,8 @@ py_test(
         "@pypi//numpy",
         "@torch_tpu//shims/torch:pytorch",
         "@torch_tpu//torch_tpu",
-        "@torch_tpu//torch_tpu/_internal/execution_mode",
-        "@torch_tpu//torch_tpu/_internal/sync",
+        "@torch_tpu//torch_tpu/src/torch_tpu/_internal/execution_mode",
+        "@torch_tpu//torch_tpu/src/torch_tpu/_internal/sync",
     ],
 )