提交 1b8689ac 编写于 作者: D DesmonDay

add code

上级 94cf213d
...@@ -188,6 +188,7 @@ DLPackTensor::DLPackTensor(phi::DenseTensor &tensor, LaneType lanes) { ...@@ -188,6 +188,7 @@ DLPackTensor::DLPackTensor(phi::DenseTensor &tensor, LaneType lanes) {
auto tensor = new DLManagedTensor; auto tensor = new DLManagedTensor;
tensor->dl_tensor = t_; tensor->dl_tensor = t_;
tensor->manager_ctx = dt_;
tensor->deleter = [](DLManagedTensor *arg) { tensor->deleter = [](DLManagedTensor *arg) {
phi::DenseTensor *tensor_ptr = phi::DenseTensor *tensor_ptr =
...@@ -200,8 +201,6 @@ DLPackTensor::DLPackTensor(phi::DenseTensor &tensor, LaneType lanes) { ...@@ -200,8 +201,6 @@ DLPackTensor::DLPackTensor(phi::DenseTensor &tensor, LaneType lanes) {
delete arg; delete arg;
}; };
tensor->manager_ctx = dt_;
return tensor; return tensor;
} }
......
...@@ -42,7 +42,6 @@ class DLPackTensor { ...@@ -42,7 +42,6 @@ class DLPackTensor {
// The shape in DLTensor is defined as int64_t* // The shape in DLTensor is defined as int64_t*
// Add this member to make TVMTensor init without heap allocation // Add this member to make TVMTensor init without heap allocation
ShapeType shape_[DDim::kMaxRank]; ShapeType shape_[DDim::kMaxRank];
phi::DenseTensor* dt_; phi::DenseTensor* dt_;
}; };
......
...@@ -473,9 +473,9 @@ void BindTensor(pybind11::module &m) { // NOLINT ...@@ -473,9 +473,9 @@ void BindTensor(pybind11::module &m) { // NOLINT
)DOC") )DOC")
.def("_to_dlpack", .def("_to_dlpack",
[](phi::DenseTensor &self) { [](phi::DenseTensor &self) {
DLPackTensor dlpack_tensor(self, 1);
pybind11::handle tensor_handle = pybind11::cast(&self); pybind11::handle tensor_handle = pybind11::cast(&self);
tensor_handle.inc_ref(); tensor_handle.inc_ref();
DLPackTensor dlpack_tensor(self, 1);
DLManagedTensor *dmt = dlpack_tensor.ToDLManagedTensor(); DLManagedTensor *dmt = dlpack_tensor.ToDLManagedTensor();
auto capsule = pybind11::capsule( auto capsule = pybind11::capsule(
static_cast<void *>(dmt), "dltensor", [](PyObject *ptr) { static_cast<void *>(dmt), "dltensor", [](PyObject *ptr) {
......
Markdown is supported
0% .
You are about to add 0 people to the discussion. Proceed with caution.
先完成此消息的编辑!
想要评论请 注册