[CustomDevice] fix set_constant (#52360)

PaddlePaddle · Mar 31, 2023 · f22b966 · f22b966
1 parent 4e23af7
commit f22b966
Showing 1 changed file with 11 additions and 2 deletions.
diff --git a/paddle/phi/kernels/funcs/math_function.cc b/paddle/phi/kernels/funcs/math_function.cc
@@ -184,7 +184,16 @@ void set_constant_with_place<phi::IPUPlace>(const phi::DeviceContext& context,
 template <>
 void set_constant_with_place<phi::CustomPlace>(
     const phi::DeviceContext& context, phi::DenseTensor* tensor, float value) {
-  PADDLE_THROW(phi::errors::Unimplemented("CustomPlace is not supported"));
+  phi::DenseTensor tmp_tensor;
+  tmp_tensor.Resize(tensor->dims());
+  context.HostAlloc(&tmp_tensor, tensor->dtype());
+  phi::VisitDataType(tmp_tensor.dtype(),
+                     TensorSetConstantCPU(&tmp_tensor, value));
+  phi::memory_utils::Copy(tensor->place(),
+                          tensor->data(),
+                          phi::CPUPlace(),
+                          tmp_tensor.data(),
+                          tensor->numel() * phi::SizeOf(tensor->dtype()));
 }
 
 template <>
@@ -230,7 +239,7 @@ void set_constant(const phi::DeviceContext& context,
   TensorSetConstantWithPlace func(context, tensor, value);
 #ifdef PADDLE_WITH_CUSTOM_DEVICE
   if (context.GetPlace().GetType() == phi::AllocationType::CUSTOM) {
-    func(phi::CPUPlace());
+    func(phi::CustomPlace());
     return;
   }
 #endif