fix (#51552)

PaddlePaddle · Mar 14, 2023 · c3f8ba9 · c3f8ba9
1 parent 2876f6f
commit c3f8ba9
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 1 deletion.
diff --git a/paddle/phi/kernels/cpu/index_select_grad_kernel.cc b/paddle/phi/kernels/cpu/index_select_grad_kernel.cc
@@ -59,5 +59,6 @@ PD_REGISTER_KERNEL(index_select_grad,
                    phi::IndexSelectGradKernel,
                    float,
                    double,
+                   phi::dtype::bfloat16,
                    int,
                    int64_t) {}
diff --git a/python/paddle/tensor/search.py b/python/paddle/tensor/search.py
@@ -350,7 +350,7 @@ def index_select(x, index, axis=0, name=None):
         check_variable_and_dtype(
             x,
             'x',
-            ['bfloat16', 'float16', 'float32', 'float64', 'int32', 'int64'],
+            ['uint16', 'float16', 'float32', 'float64', 'int32', 'int64'],
             'paddle.tensor.search.index_select',
         )
         check_variable_and_dtype(