diff --git a/paddle/phi/api/lib/scalar.cc b/paddle/phi/api/lib/scalar.cc index c31338de09f1e2864a8920a2297de9446b8bce29..09b15f629d8040784c78918d109d85e5fa69b277 100644 --- a/paddle/phi/api/lib/scalar.cc +++ b/paddle/phi/api/lib/scalar.cc @@ -37,6 +37,12 @@ ScalarBase::ScalarBase(const Tensor& tensor_in) GetDataFromTensor(dst_tensor); } else if (tensor_in_place == phi::AllocationType::CPU) { GetDataFromTensor(tensor_in); +#ifdef PADDLE_WITH_CUSTOM_DEVICE + } else if (tensor_in_place == phi::AllocationType::CUSTOM) { + Tensor dst_tensor; + copy(tensor_in, phi::CPUPlace(), true, &dst_tensor); + GetDataFromTensor(dst_tensor); +#endif } else { PADDLE_THROW(phi::errors::Unimplemented( "Now, it is not supported to construct Scalar using tensor that its " diff --git a/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py b/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py index bf1effe21917f55928225cf6a57961a3a99a5135..07e225160407fe7aa1048e7b87c9fdee964c51fc 100644 --- a/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py +++ b/python/paddle/fluid/tests/custom_runtime/test_custom_cpu_plugin.py @@ -41,6 +41,7 @@ class TestCustomCPUPlugin(unittest.TestCase): self._test_eager_backward_api() self._test_eager_copy_to() self._test_fallback_kernel() + self._test_scalar() self._test_custom_device_dataloader() self._test_custom_device_mnist() @@ -170,6 +171,13 @@ class TestCustomCPUPlugin(unittest.TestCase): z = paddle.add(x, y) np.testing.assert_array_equal(z, r) + def _test_scalar(self): + import paddle + data_1 = paddle.to_tensor([[[[1.0, 4.0, 5.0, 7.0], [3.0, 4.0, 5.0, + 6.0]]]]) + k_t = paddle.to_tensor([3], dtype="int32") + value_1, indices_1 = paddle.topk(data_1, k=k_t) + def tearDown(self): del os.environ['CUSTOM_DEVICE_ROOT']