add bf16 for numel (#48121)

a7d306af · Yuang Liu · GitHub · 9918bf9c · a7d306af · a7d306af
隐藏空白更改
内联并排

Showing with 2 addition and 0 deletion

paddle/phi/kernels/cpu/size_kernel.cc paddle/phi/kernels/cpu/size_kernel.cc +1 -0

paddle/phi/kernels/gpu/size_kernel.cu paddle/phi/kernels/gpu/size_kernel.cu +1 -0

未找到文件。
--- a/paddle/phi/kernels/cpu/size_kernel.cc
+++ b/paddle/phi/kernels/cpu/size_kernel.cc
@@ -27,6 +27,7 @@ PD_REGISTER_KERNEL(size,
                   int,
                   int64_t,
                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
                   float,
                   double,
                   bool) {}
--- a/paddle/phi/kernels/gpu/size_kernel.cu
+++ b/paddle/phi/kernels/gpu/size_kernel.cu
@@ -26,6 +26,7 @@ PD_REGISTER_KERNEL(size,
                   int,
                   int64_t,
                   phi::dtype::float16,
+                   phi::dtype::bfloat16,
                   float,
                   double,
                   bool) {}