fix the test

openxla · Sep 12, 2024 · b5e43d6 · b5e43d6
1 parent 8113e6f
commit b5e43d6
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc b/xla/pjrt/gpu/se_gpu_pjrt_client_test.cc
@@ -1236,7 +1236,7 @@ TEST(StreamExecutorGpuClientTest,
 
   auto device = client->addressable_devices()[0];
   TF_EXPECT_OK(device->default_memory_space());
-  TF_ASSIGN_OR_RETURN(
+  TF_ASSERT_OK_AND_ASSIGN(
       auto input, client->BufferFromHostBuffer(
                       data.data(), shape.element_type(), shape.dimensions(),
                       /*byte_strides=*/std::nullopt,