Skip to content

Commit dcd4e09

Browse files
committed
Fix memory leaks
Add missing wrapping of raw pointers in smart pointers, so that the destructors of the underlying TensorRT objects are called properly Signed-off-by: Gabriel Cuendet <[email protected]>
1 parent f0e6d2d commit dcd4e09

File tree

2 files changed

+2
-2
lines changed

2 files changed

+2
-2
lines changed

core/conversion/conversionctx/ConversionCtx.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -164,7 +164,7 @@ void ConversionCtx::RecordNewITensor(const torch::jit::Value* value, nvinfer1::I
164164

165165
std::string ConversionCtx::SerializeEngine() {
166166
#if NV_TENSORRT_MAJOR > 7
167-
auto serialized_network = builder->buildSerializedNetwork(*net, *cfg);
167+
auto serialized_network = make_trt(builder->buildSerializedNetwork(*net, *cfg));
168168
if (!serialized_network) {
169169
TORCHTRT_THROW_ERROR("Building serialized network failed in TensorRT");
170170
}

core/runtime/register_jit_hooks.cpp

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -87,7 +87,7 @@ static auto TORCHTRT_UNUSED TRTEngineTSRegistrtion =
8787
.def_pickle(
8888
[](const c10::intrusive_ptr<TRTEngine>& self) -> std::vector<std::string> {
8989
// Serialize TensorRT engine
90-
auto serialized_trt_engine = self->cuda_engine->serialize();
90+
auto serialized_trt_engine = make_trt(self->cuda_engine->serialize());
9191

9292
// Adding device info related meta data to the serialized file
9393
auto trt_engine = std::string((const char*)serialized_trt_engine->data(), serialized_trt_engine->size());

0 commit comments

Comments
 (0)