From 43b950f799551ea94a7bbd9a3ac166b87ffc1a84 Mon Sep 17 00:00:00 2001
From: YuhangLi <104877312+piDack@users.noreply.github.com>
Date: Fri, 21 Apr 2023 09:49:12 +0800
Subject: [PATCH] [CUSTOM_DEVICE]custom device fetch_v2 support (#53093)

---
 paddle/fluid/operators/controlflow/fetch_v2_op.cc | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/paddle/fluid/operators/controlflow/fetch_v2_op.cc b/paddle/fluid/operators/controlflow/fetch_v2_op.cc
index 48065cc3fd1..9105fe88862 100644
--- a/paddle/fluid/operators/controlflow/fetch_v2_op.cc
+++ b/paddle/fluid/operators/controlflow/fetch_v2_op.cc
@@ -156,7 +156,8 @@ class FetchV2Kernel {
       }
       auto *dst_item = &(PADDLE_GET(phi::DenseTensor, fetch_list->at(col)));
       bool check_place = platform::is_cpu_place(src_item.place()) ||
-                         platform::is_cuda_pinned_place(src_item.place());
+                         platform::is_cuda_pinned_place(src_item.place()) ||
+                         platform::is_custom_place(src_item.place());
       PADDLE_ENFORCE_EQ(
           check_place,
           true,
-- 
GitLab