[CUSTOM_DEVICE]custom device fetch_v2 support (#53093)

43b950f7 · YuhangLi · GitHub · 63c83870 · 43b950f7
隐藏空白更改
内联并排

Showing with 2 addition and 1 deletion

paddle/fluid/operators/controlflow/fetch_v2_op.cc paddle/fluid/operators/controlflow/fetch_v2_op.cc +2 -1

未找到文件。
--- 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,