diff --git a/reverb/cc/BUILD b/reverb/cc/BUILD
index d41d689..23cd3dd 100644
--- a/reverb/cc/BUILD
+++ b/reverb/cc/BUILD
@@ -194,7 +194,7 @@ reverb_cc_library(
         "//third_party/absl/status",
         "//third_party/absl/status:statusor",
         "//third_party/absl/strings",
-        "//third_party/tensorflow/tsl/platform:status",
+        "//third_party/tensorflow/compiler/xla/tsl/platform:status",
     ] + reverb_tf_deps(),
 )
 
diff --git a/reverb/cc/tensor_compression.cc b/reverb/cc/tensor_compression.cc
index 5cab855..31be10e 100644
--- a/reverb/cc/tensor_compression.cc
+++ b/reverb/cc/tensor_compression.cc
@@ -21,11 +21,11 @@
 #include "absl/strings/str_cat.h"
 #include "reverb/cc/platform/logging.h"
 #include "reverb/cc/platform/snappy.h"
+#include "tensorflow/compiler/xla/tsl/platform/status.h"
 #include "tensorflow/core/framework/register_types.h"
 #include "tensorflow/core/framework/tensor_shape.h"
 #include "tensorflow/core/framework/types.h"
 #include "tensorflow/core/framework/types.pb.h"
-#include "tensorflow/tsl/platform/status.h"
 
 namespace deepmind {
 namespace reverb {