feat(api): log CUDA memory stats (#141)
This commit is contained in:
parent
03a579a535
commit
005650a9a2
|
@ -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,
|
||||
)
|
||||
|
|
Loading…
Reference in New Issue