fix dlpack memory leak (InternLM#1344)

irexyc · Mar 26, 2024 · 9c954ef · 9c954ef
1 parent 893a574
commit 9c954ef
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/turbomind/python/bind.cpp b/src/turbomind/python/bind.cpp
@@ -306,7 +306,7 @@ PYBIND11_MODULE(_turbomind, m)
             "__dlpack__",
             [](triton::Tensor* self, long stream) {
                 auto tensor_ptr = TritonTensorToDLManagedTensor(*self);
-                return new py::capsule(tensor_ptr.release(), kDlTensorCapsuleName, [](PyObject* obj) {
+                return py::capsule(tensor_ptr.release(), kDlTensorCapsuleName, [](PyObject* obj) {
                     DLManagedTensor* dlmt =
                         static_cast<DLManagedTensor*>(PyCapsule_GetPointer(obj, kDlTensorCapsuleName));
                     if (dlmt) {