From 5bd8a018b7c372548b8ecbd7a4c4bd982b9055eb Mon Sep 17 00:00:00 2001
From: Kyle Lucke <klucke@google.com>
Date: Thu, 2 May 2024 11:32:57 -0700
Subject: [PATCH] Replace one last use of ScopedDeviceMemory in
 buffer_comparator_test with DeviceHandle.

PiperOrigin-RevId: 630128405
---
 .../xla/xla/service/gpu/buffer_comparator_test.cc    | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/third_party/xla/xla/service/gpu/buffer_comparator_test.cc b/third_party/xla/xla/service/gpu/buffer_comparator_test.cc
index 5c1c8d92022711..9faa3156a84965 100644
--- a/third_party/xla/xla/service/gpu/buffer_comparator_test.cc
+++ b/third_party/xla/xla/service/gpu/buffer_comparator_test.cc
@@ -351,20 +351,20 @@ TEST_F(BufferComparatorTest, BF16) {
 
   auto stream = stream_exec_->CreateStream().value();
 
-  se::ScopedDeviceMemory<Eigen::bfloat16> lhs(
+  se::DeviceMemoryHandle lhs(
       stream_exec_,
       stream_exec_->AllocateArray<Eigen::bfloat16>(element_count));
-  InitializeBuffer(stream.get(), BF16, &rng_state, *lhs.ptr());
+  InitializeBuffer(stream.get(), BF16, &rng_state, lhs.memory());
 
-  se::ScopedDeviceMemory<Eigen::bfloat16> rhs(
+  se::DeviceMemoryHandle rhs(
       stream_exec_,
       stream_exec_->AllocateArray<Eigen::bfloat16>(element_count));
-  InitializeBuffer(stream.get(), BF16, &rng_state, *rhs.ptr());
+  InitializeBuffer(stream.get(), BF16, &rng_state, rhs.memory());
 
   BufferComparator comparator(ShapeUtil::MakeShape(BF16, {element_count}),
                               HloModuleConfig());
-  EXPECT_FALSE(
-      comparator.CompareEqual(stream.get(), *lhs.ptr(), *rhs.ptr()).value());
+  EXPECT_FALSE(comparator.CompareEqual(stream.get(), lhs.memory(), rhs.memory())
+                   .value());
 }
 
 }  // namespace