From d7d490e413a89e8595cf50d7547c7098f6334d59 Mon Sep 17 00:00:00 2001
From: DesmonDay <908660116@qq.com>
Date: Wed, 1 Feb 2023 07:01:31 +0000
Subject: [PATCH] fix to_dlpack for loop

---
 paddle/fluid/pybind/tensor.cc      | 18 ++++++------------
 python/paddle/tests/test_dlpack.py |  6 ++++++
 2 files changed, 12 insertions(+), 12 deletions(-)

diff --git a/paddle/fluid/pybind/tensor.cc b/paddle/fluid/pybind/tensor.cc
index 4bdde24f431..29a17d25d9d 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<void *>(dmt), "dltensor", [](PyObject *ptr) {
-                   if (ptr) {
-                     auto dltensor = new DLManagedTensor;
-                     try {
-                       dltensor = reinterpret_cast<DLManagedTensor *>(
-                           PyCapsule_GetPointer(ptr, "used_dltensor"));
-                       return;
-                     } catch (...) {
-                       dltensor = reinterpret_cast<DLManagedTensor *>(
-                           PyCapsule_GetPointer(ptr, "dltensor"));
-                     }
-                     dltensor->deleter(dltensor);
+                   if (!PyCapsule_IsValid(ptr, "dltensor")) {
+                     return;
                    }
+                   DLManagedTensor *dmt = static_cast<DLManagedTensor *>(
+                       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 77ffdbecedb..504ff4c48c6 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):
-- 
GitLab