From a0bbfbd498823409a5b07c0d4f2716008eea5481 Mon Sep 17 00:00:00 2001 From: Aganlengzi Date: Tue, 16 Aug 2022 17:20:40 +0800 Subject: [PATCH] support fp16 softmax on custom place (#45177) --- paddle/fluid/operators/softmax_op.cc | 10 ++++++---- 1 file changed, 6 insertions(+), 4 deletions(-) diff --git a/paddle/fluid/operators/softmax_op.cc b/paddle/fluid/operators/softmax_op.cc index 3d85c9b0a4..d468e4a17f 100644 --- a/paddle/fluid/operators/softmax_op.cc +++ b/paddle/fluid/operators/softmax_op.cc @@ -62,10 +62,11 @@ class SoftmaxOp : public framework::OperatorWithKernel { platform::is_gpu_place(ctx.GetPlace()) || platform::is_npu_place(ctx.GetPlace()) || platform::is_xpu_place(ctx.GetPlace()) || - platform::is_mlu_place(ctx.GetPlace()), + platform::is_mlu_place(ctx.GetPlace()) || + platform::is_custom_place(ctx.GetPlace()), true, platform::errors::InvalidArgument( - "float16 can only be used on GPU/NPU/XPU/MLU place")); + "float16 can only be used on GPU/NPU/XPU/MLU and custom place")); } return framework::OpKernelType( @@ -176,9 +177,10 @@ class SoftmaxOpGrad : public framework::OperatorWithKernel { if (!(platform::is_gpu_place(ctx.GetPlace()) || platform::is_npu_place(ctx.GetPlace()) || platform::is_xpu_place(ctx.GetPlace()) || - platform::is_mlu_place(ctx.GetPlace()))) + platform::is_mlu_place(ctx.GetPlace()) || + platform::is_custom_place(ctx.GetPlace()))) PADDLE_THROW(platform::errors::InvalidArgument( - "float16 can only be used on GPU/NPU/XPU/MLU place")); + "float16 can only be used on GPU/NPU/XPU/MLU and custom place")); } return framework::OpKernelType( -- GitLab