未验证 提交 f55eb06f 编写于 作者: H houj04 提交者: GitHub

[XPU] fix unittest of shape op. (#54323)

上级 e7a38f15
......@@ -16,6 +16,7 @@ limitations under the License. */
#include "paddle/phi/backends/cpu/cpu_context.h"
#include "paddle/phi/backends/gpu/gpu_context.h"
#include "paddle/phi/backends/xpu/xpu_context.h"
#include "paddle/phi/common/float16.h"
#include "paddle/phi/core/kernel_registry.h"
#include "paddle/phi/kernels/shape_kernel.h"
......@@ -71,6 +72,23 @@ PD_REGISTER_KERNEL(shape_sr,
}
#endif
#if defined(PADDLE_WITH_XPU)
PD_REGISTER_KERNEL(shape_sr,
XPU,
ALL_LAYOUT,
phi::sr::ShapeKernel,
bool,
int,
int64_t,
float,
double,
phi::dtype::float16) {
kernel->InputAt(0).SetBackend(phi::Backend::ALL_BACKEND);
kernel->OutputAt(0).SetBackend(phi::Backend::CPU);
kernel->OutputAt(0).SetDataType(phi::DataType::INT32);
}
#endif
#ifdef PADDLE_WITH_CUSTOM_DEVICE
PD_REGISTER_KERNEL(shape_sr,
Custom,
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册