From 005650a9a228093e44f22fc9868d8804c810b173 Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Thu, 16 Feb 2023 23:51:17 -0600 Subject: [PATCH] feat(api): log CUDA memory stats (#141) --- api/onnx_web/utils.py | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py index 733326e5..a6f35a34 100644 --- a/api/onnx_web/utils.py +++ b/api/onnx_web/utils.py @@ -144,3 +144,9 @@ def run_gc(devices: List[DeviceParams] = []): with torch.cuda.device(device.torch_str()): torch.cuda.empty_cache() torch.cuda.ipc_collect() + mem_free, mem_total = torch.cuda.mem_get_info(device=device.torch_str()) + logger.debug( + "remaining CUDA VRAM usage: %s of %s", + (mem_total - mem_free), + mem_total, + )