diff --git a/paddle/fluid/operators/memcpy_op.cc b/paddle/fluid/operators/memcpy_op.cc index ef430f8bfaf4368d34a0f447157864f4fb4c793f..9fb06c5968cbdaa2df471c831eac3a9e57884df5 100644 --- a/paddle/fluid/operators/memcpy_op.cc +++ b/paddle/fluid/operators/memcpy_op.cc @@ -113,7 +113,9 @@ class MemcpyOpProtoMaker : public framework::OpProtoAndCheckerMaker { "1: dst is on CUDAPlace. " "2: dst is on CUDAPinnedPlace. " "3: dst is on XPUPlace. " - "4: dst is on NPUPlace. "); + "4: dst is on NPUPlace. " + "5: dst is on NPUPinnerPlace. " + "6: dst is on CustomDevicePlace"); AddComment(R"DOC( Memcpy Operator. By now, it ONLY supports the memcopy between CUDAPinnedPlace <-> CUDAPlace or diff --git a/paddle/fluid/operators/memcpy_op.h b/paddle/fluid/operators/memcpy_op.h index a35fefa53ba3d4f206f26b5ee84afbb9573015f7..092d0cf368e6a6e019164adcf71536459f39f30f 100644 --- a/paddle/fluid/operators/memcpy_op.h +++ b/paddle/fluid/operators/memcpy_op.h @@ -41,6 +41,7 @@ class MemcpyFunctor { XPU = 3, NPU = 4, NPU_PINNED = 5, + CUSTOM_DEVICE = 6, }; public: @@ -67,6 +68,11 @@ class MemcpyFunctor { } else if (dst_place_type_ == DeviceType::NPU_PINNED) { /* npu->npu_pin */ framework::TensorCopy( lod_tensor, platform::NPUPinnedPlace(), dev_ctx_, &out_tensor); +#endif +#ifdef PADDLE_WTIH_CUSTOM_DEVICE + } else if (dst_place_type_ == DeviceType::CUSTOM_DEVICE) { + framework::TensorCopy( + lod_tensor, dev_ctx_.GetPlace(), dev_ctx_, &out_tensor); #endif } else { PADDLE_THROW(platform::errors::Unimplemented(