Merge pull request #376 from NVIDIA/fix_seralization_memory_leak

fix(//core/conversion/conversionctx): Fix memory leak in conversionctx

Author: narendasan

core/conversion/conversionctx/ConversionCtx.cpp

@@ -148,7 +148,9 @@ std::string ConversionCtx::SerializeEngine() {
   auto engine = builder->buildEngineWithConfig(*net, *cfg);
   auto serialized_engine = engine->serialize();
   engine->destroy();
-  return std::string((const char*)serialized_engine->data(), serialized_engine->size());
+  auto engine_str = std::string((const char*)serialized_engine->data(), serialized_engine->size());
+  serialized_engine->destroy();
+  return engine_str;
 }
 
 bool ConversionCtx::CheckLayerAddition(const torch::jit::Node* n) {

py/trtorch/_compiler.py

@@ -157,5 +157,6 @@ def get_build_info() -> str:
     build_info = "TRTorch Version: " + str(__version__) + '\n' + build_info
     return build_info
 
+
 def set_device(gpu_id):
     trtorch._C.set_device(gpu_id)

tests/py/test_multi_gpu.py

@@ -5,7 +5,9 @@
 
 from model_test_case import ModelTestCase
 
+
 class TestMultiGpuSwitching(ModelTestCase):
+
     def setUp(self):
         if torch.cuda.device_count() < 2:
             self.fail("Test is not relevant for this platform since number of available CUDA devices is less than 2")
@@ -55,12 +57,14 @@ def test_compile_script(self):
         trtorch.set_device(0)
         self.assertTrue(same < 2e-3)
 
+
 def test_suite():
     suite = unittest.TestSuite()
     suite.addTest(TestMultiGpuSwitching.parametrize(TestMultiGpuSwitching, model=models.resnet18(pretrained=True)))
 
     return suite
 
+
 suite = test_suite()
 
 runner = unittest.TextTestRunner()