From f596714c5962023ac1fab07b0bc8e25e461b8f54 Mon Sep 17 00:00:00 2001 From: elaheh Date: Thu, 29 Jan 2026 09:55:12 -0500 Subject: [PATCH] Fix memory leak by adding py::return_value_policy::move The pybind11 bindings were missing an explicit return value policy, causing tensor references to not be properly released. This resulted in ~6MB leaked per forward call, matching the output tensor size. Adding py::return_value_policy::move transfers full ownership to Python, allowing proper garbage collection. Test results: - Before: 2377 MB leaked over 200 iterations - After: 4 MB stable (no leak) --- ppisp/ext.cpp | 7 +++++-- 1 file changed, 5 insertions(+), 2 deletions(-) diff --git a/ppisp/ext.cpp b/ppisp/ext.cpp index 02c5a0e..0e29a2e 100644 --- a/ppisp/ext.cpp +++ b/ppisp/ext.cpp @@ -18,6 +18,9 @@ #include "bindings.h" PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { - m.def("ppisp_forward", &ppisp_forward_tensor); - m.def("ppisp_backward", &ppisp_backward_tensor); + // Use move policy to ensure proper tensor ownership transfer + m.def("ppisp_forward", &ppisp_forward_tensor, + py::return_value_policy::move); + m.def("ppisp_backward", &ppisp_backward_tensor, + py::return_value_policy::move); }