pylayer support HIP (#38184)

PaddlePaddle · Dec 16, 2021 · 2e76d5a · 2e76d5a
1 parent af30f54
commit 2e76d5a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/paddle/fluid/operators/py_layer_op.cc b/paddle/fluid/operators/py_layer_op.cc
@@ -220,7 +220,7 @@ REGISTER_OP_CPU_KERNEL(
                          ::paddle::platform::complex<float>>,
     ops::PyLayerOpKernel<paddle::platform::CPUDeviceContext,
                          ::paddle::platform::complex<double>>);
-#ifdef PADDLE_WITH_CUDA
+#if defined(PADDLE_WITH_CUDA) || defined(PADDLE_WITH_HIP)
 REGISTER_OP_CUDA_KERNEL(
     py_layer, ops::PyLayerOpKernel<paddle::platform::CUDADeviceContext, float>,
     ops::PyLayerOpKernel<paddle::platform::CUDADeviceContext,
@@ -239,4 +239,4 @@ REGISTER_OP_CUDA_KERNEL(
                          ::paddle::platform::complex<float>>,
     ops::PyLayerOpKernel<paddle::platform::CUDADeviceContext,
                          ::paddle::platform::complex<double>>);
-#endif  // PADDLE_WITH_CUDA
+#endif  // PADDLE_WITH_CUDA || PADDLE_WITH_HIP