diff --git a/python/src/main.cpp b/python/src/main.cpp index 1fd9473d3..10b28f5b6 100644 --- a/python/src/main.cpp +++ b/python/src/main.cpp @@ -884,7 +884,7 @@ PYBIND11_MODULE(ncnn, m) #endif // NCNN_STDIO .def("clear", &Net::clear) - .def("create_extractor", &Net::create_extractor) + .def("create_extractor", &Net::create_extractor, py::keep_alive<0, 1>()) //net should be kept alive until retuned ex is freed by gc .def("blobs", &Net::blobs, py::return_value_policy::reference_internal) .def("layers", &Net::layers, py::return_value_policy::reference_internal); @@ -1113,8 +1113,8 @@ PYBIND11_MODULE(ncnn, m) m.def("destroy_gpu_instance", &destroy_gpu_instance); m.def("get_gpu_count", &get_gpu_count); m.def("get_default_gpu_index", &get_default_gpu_index); - m.def("get_gpu_info", &get_gpu_info, py::arg("device_index") = 0, py::return_value_policy::reference_internal); - m.def("get_gpu_device", &get_gpu_device, py::arg("device_index") = 0, py::return_value_policy::reference_internal); + m.def("get_gpu_info", &get_gpu_info, py::arg("device_index") = 0, py::return_value_policy::reference); + m.def("get_gpu_device", &get_gpu_device, py::arg("device_index") = 0, py::return_value_policy::reference); py::class_(m, "VkBufferMemory") .def_readwrite("offset", &VkBufferMemory::offset) diff --git a/python/tests/benchmark.py b/python/tests/benchmark.py index 21e5804ba..6f0bbfdf7 100644 --- a/python/tests/benchmark.py +++ b/python/tests/benchmark.py @@ -56,6 +56,7 @@ def benchmark(comment, _in, opt): # warm up for i in range(g_warmup_loop_count): + # test with statement with net.create_extractor() as ex: ex.input("data", _in) ex.extract("output") @@ -67,9 +68,10 @@ def benchmark(comment, _in, opt): for i in range(g_loop_count): start = time.time() - with net.create_extractor() as ex: - ex.input("data", _in) - ex.extract("output") + # test net keep alive until ex freed + ex = net.create_extractor() + ex.input("data", _in) + ex.extract("output") end = time.time()