From 4a77cf53796f20608ae334bfc264f933b243fd5f Mon Sep 17 00:00:00 2001
From: houj04 <35131887+houj04@users.noreply.github.com>
Date: Fri, 9 Jun 2023 09:50:10 +0800
Subject: [PATCH] [XPU] add registration of SplitWithNumKernel with int64.
 (#54478)

---
 paddle/phi/backends/xpu/xpu2_op_list.cc | 3 ++-
 paddle/phi/kernels/xpu/split_kernel.cc  | 1 +
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/paddle/phi/backends/xpu/xpu2_op_list.cc b/paddle/phi/backends/xpu/xpu2_op_list.cc
index 7a287c99339..8b9a9fba639 100644
--- a/paddle/phi/backends/xpu/xpu2_op_list.cc
+++ b/paddle/phi/backends/xpu/xpu2_op_list.cc
@@ -726,7 +726,8 @@ XPUOpMap& get_kl2_ops() {
       {"split_with_num",
        XPUKernelSet({phi::DataType::FLOAT32,
                      phi::DataType::FLOAT16,
-                     phi::DataType::INT32})},
+                     phi::DataType::INT32,
+                     phi::DataType::INT64})},
       {"sqrt", XPUKernelSet({phi::DataType::FLOAT32, phi::DataType::FLOAT16})},
       {"sqrt_grad", XPUKernelSet({phi::DataType::FLOAT32})},
       {"square_grad",
diff --git a/paddle/phi/kernels/xpu/split_kernel.cc b/paddle/phi/kernels/xpu/split_kernel.cc
index f94dbaac037..410ca7e1fe6 100644
--- a/paddle/phi/kernels/xpu/split_kernel.cc
+++ b/paddle/phi/kernels/xpu/split_kernel.cc
@@ -79,5 +79,6 @@ PD_REGISTER_KERNEL(split_with_num,
                    ALL_LAYOUT,
                    phi::SplitWithNumKernel,
                    float,
+                   int64_t,
                    int,
                    phi::dtype::float16) {}
-- 
GitLab