diff --git a/paddle/fluid/pybind/tensor.cc b/paddle/fluid/pybind/tensor.cc index 4bdde24f431bc830171ea4d3c02d1064da67926b..29a17d25d9dc24c29cecf4c57f27fd45c8d646c2 100644 --- a/paddle/fluid/pybind/tensor.cc +++ b/paddle/fluid/pybind/tensor.cc @@ -475,20 +475,14 @@ void BindTensor(pybind11::module &m) { // NOLINT [](phi::DenseTensor &self) { DLPackTensor dlpack_tensor(self, 1); DLManagedTensor *dmt = dlpack_tensor.ToDLManagedTensor(); - auto capsule = py::capsule( + auto capsule = pybind11::capsule( static_cast(dmt), "dltensor", [](PyObject *ptr) { - if (ptr) { - auto dltensor = new DLManagedTensor; - try { - dltensor = reinterpret_cast( - PyCapsule_GetPointer(ptr, "used_dltensor")); - return; - } catch (...) { - dltensor = reinterpret_cast( - PyCapsule_GetPointer(ptr, "dltensor")); - } - dltensor->deleter(dltensor); + if (!PyCapsule_IsValid(ptr, "dltensor")) { + return; } + DLManagedTensor *dmt = static_cast( + PyCapsule_GetPointer(ptr, "dltensor")); + dmt->deleter(dmt); }); return capsule; }) diff --git a/python/paddle/tests/test_dlpack.py b/python/paddle/tests/test_dlpack.py index 77ffdbecedbf60317f5209b4be72bf477b9f10a8..504ff4c48c66eda9987ebf44ebb9dc794b18bf1f 100644 --- a/python/paddle/tests/test_dlpack.py +++ b/python/paddle/tests/test_dlpack.py @@ -116,6 +116,12 @@ class TestDLPack(unittest.TestCase): dlpack = paddle.utils.dlpack.to_dlpack(a) b = paddle.utils.dlpack.from_dlpack(dlpack) + def test_to_dlpack_for_loop(self): + # See Paddle issue 50120 + for i in range(10): + x = paddle.rand([3, 5]) + dlpack = paddle.utils.dlpack.to_dlpack(x) + class TestRaiseError(unittest.TestCase): def test_from_dlpack_raise_type_error(self):