mirrors/mmdeploy
1
0
Fork 0
mirror of https://github.com/open-mmlab/mmdeploy.git synced 2025-01-14 08:09:43 +08:00
Code Issues Projects Releases Wiki Activity

[Fix] fix torch allocator resouce releasing (#1708)

Browse Source 
* delete root logger and add condition before calling caching_allocator_delete

* fix lint error

* use torch._C._cuda_cudaCachingAllocator_raw_delete
This commit is contained in:
AllentDan 2023-02-06 11:35:44 +08:00 committed by GitHub
parent b85f34141b
commit 12b3d18c7a
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 2 additions and 3 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

5
mmdeploy/backend/tensorrt/torch_allocator.py
View File

@ -13,6 +13,7 @@ class TorchAllocator(trt.IGpuAllocator):
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 @@ class TorchAllocator(trt.IGpuAllocator):
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