diff --git a/paddle/fluid/framework/dlpack_tensor.cc b/paddle/fluid/framework/dlpack_tensor.cc index 419eb0e7fe849c1890774c075fa9fcd0a55e4fab..e1131e4c6d994f9cdc678a560593f794cc1d2c42 100644 --- a/paddle/fluid/framework/dlpack_tensor.cc +++ b/paddle/fluid/framework/dlpack_tensor.cc @@ -188,6 +188,7 @@ DLPackTensor::DLPackTensor(phi::DenseTensor &tensor, LaneType lanes) { auto tensor = new DLManagedTensor; tensor->dl_tensor = t_; + tensor->manager_ctx = dt_; tensor->deleter = [](DLManagedTensor *arg) { phi::DenseTensor *tensor_ptr = @@ -200,8 +201,6 @@ DLPackTensor::DLPackTensor(phi::DenseTensor &tensor, LaneType lanes) { delete arg; }; - tensor->manager_ctx = dt_; - return tensor; } diff --git a/paddle/fluid/framework/dlpack_tensor.h b/paddle/fluid/framework/dlpack_tensor.h index 2fcd2511afb8b5122ec81fc5ac0cf4d7f1867b9a..10c38a956cfda5d3b24919a73b966f778f653531 100644 --- a/paddle/fluid/framework/dlpack_tensor.h +++ b/paddle/fluid/framework/dlpack_tensor.h @@ -42,7 +42,6 @@ class DLPackTensor { // The shape in DLTensor is defined as int64_t* // Add this member to make TVMTensor init without heap allocation ShapeType shape_[DDim::kMaxRank]; - phi::DenseTensor* dt_; }; diff --git a/paddle/fluid/pybind/tensor.cc b/paddle/fluid/pybind/tensor.cc index b58e697e69532b0083cad77b3aedba0236449e32..8b020792b6eb2205f2e7da7bdc6a32b03115ca0c 100644 --- a/paddle/fluid/pybind/tensor.cc +++ b/paddle/fluid/pybind/tensor.cc @@ -473,9 +473,9 @@ void BindTensor(pybind11::module &m) { // NOLINT )DOC") .def("_to_dlpack", [](phi::DenseTensor &self) { - DLPackTensor dlpack_tensor(self, 1); pybind11::handle tensor_handle = pybind11::cast(&self); tensor_handle.inc_ref(); + DLPackTensor dlpack_tensor(self, 1); DLManagedTensor *dmt = dlpack_tensor.ToDLManagedTensor(); auto capsule = pybind11::capsule( static_cast(dmt), "dltensor", [](PyObject *ptr) {