diff --git a/test/custom_runtime/custom_op.cc b/test/custom_runtime/custom_op.cc
index 68c06dcd94be262cf7b8aaa7680586bc1afa7913..5ae35d8bda44892ecd26e6989c635cf68b01ddd8 100644
--- a/test/custom_runtime/custom_op.cc
+++ b/test/custom_runtime/custom_op.cc
@@ -199,7 +199,7 @@ std::vector<paddle::Tensor> StreamForward(const paddle::Tensor& x) {
   auto dev_ctx =
       paddle::experimental::DeviceContextPool::Instance().Get(x.place());
   auto custom_ctx = static_cast<const phi::CustomContext*>(dev_ctx);
-  void* stream = custom_ctx->stream();
+  std::shared_ptr<phi::stream::Stream> stream = custom_ctx->GetStream();
 
   PD_CHECK(stream != nullptr);
   std::cout << "Check stream != nullptr successfully" << std::endl;