diff --git a/paddle/phi/kernels/funcs/math_function.cc b/paddle/phi/kernels/funcs/math_function.cc
index 84df3c4b139aa4..f4c6d2db94e90c 100644
--- a/paddle/phi/kernels/funcs/math_function.cc
+++ b/paddle/phi/kernels/funcs/math_function.cc
@@ -71,7 +71,6 @@ template struct SetConstant<phi::XPUContext, int64_t>;
 template struct SetConstant<phi::XPUContext, bool>;
 template struct SetConstant<phi::XPUContext, phi::dtype::complex<float>>;
 template struct SetConstant<phi::XPUContext, phi::dtype::complex<double>>;
-
 #endif
 
 #define DEFINE_CPU_TRANS(RANK)                                            \
@@ -246,7 +245,12 @@ void set_constant(const phi::DeviceContext& context,
   // tensor->place().apply_visitor(func);
   phi::VisitPlace(tensor->place(), func);
 #elif defined(PADDLE_WITH_XPU)
-  func(phi::XPUPlace());
+  if (context.GetPlace().GetType() == phi::AllocationType::XPU) {
+    func(phi::XPUPlace());
+    return;
+  } else {
+    func(phi::CPUPlace());
+  }
 #else
   func(phi::CPUPlace());
 #endif
diff --git a/paddle/phi/kernels/funcs/math_function.h b/paddle/phi/kernels/funcs/math_function.h
index 5e5834bf91e307..2e1013ec7fc1be 100644
--- a/paddle/phi/kernels/funcs/math_function.h
+++ b/paddle/phi/kernels/funcs/math_function.h
@@ -143,9 +143,9 @@ struct TensorSetConstantXPU<float> {
     auto* ctx = phi::DeviceContextPool::Instance().Get(place_);
     auto begin = ctx->Alloc<T>(tensor_);
     int numel = tensor_->numel();
-    if (((std::is_same<T, float>::value) ||
-         (std::is_same<T, phi::dtype::float16>::value)) &&
-        (place_ == phi::XPUPlace())) {
+    if ((std::is_same<T, float>::value) ||
+        (std::is_same<T, phi::dtype::bfloat16>::value) ||
+        (std::is_same<T, phi::dtype::float16>::value)) {
       using XPUType = typename XPUTypeTrait<T>::Type;
       auto* dev_ctx = static_cast<phi::XPUContext*>(ctx);
       int r = xpu::constant<XPUType>(dev_ctx->x_context(),