fix torch allocator resouce releasing (#1711)

open-mmlab · Feb 6, 2023 · a9f8d8c · a9f8d8c
1 parent 87e4814
commit a9f8d8c
Showing 1 changed file with 2 additions and 3 deletions.
diff --git a/mmdeploy/backend/tensorrt/torch_allocator.py b/mmdeploy/backend/tensorrt/torch_allocator.py
@@ -13,6 +13,7 @@ def __init__(self, device_id: int = 0) -> None:
 
         self.device_id = device_id
         self.mems = set()
+        self.caching_delete = torch._C._cuda_cudaCachingAllocator_raw_delete
 
     def __del__(self):
         """destructor."""
@@ -53,11 +54,9 @@ def deallocate(self: trt.IGpuAllocator, memory: int) -> bool:
         Returns:
             bool: deallocate success.
         """
-        logger = get_root_logger()
-        logger.debug(f'deallocate {memory} with TorchAllocator.')
         if memory not in self.mems:
             return False
 
-        torch.cuda.caching_allocator_delete(memory)
+        self.caching_delete(memory)
         self.mems.discard(memory)
         return True