apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): correct device for VRAM check

This commit is contained in:
Sean Sube 2023-02-17 08:44:42 -06:00
parent 56a826228d
commit cf8ca34845
Signed by: ssube
GPG Key ID: 3EED7B957D362AF1
1 changed files with 1 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
api/onnx_web/utils.py
View File

@ -144,7 +144,7 @@ def run_gc(devices: List[DeviceParams] = []):
with torch.cuda.device(device.torch_str()): with torch.cuda.device(device.torch_str()):
torch.cuda.empty_cache() torch.cuda.empty_cache()
torch.cuda.ipc_collect() torch.cuda.ipc_collect()
mem_free, mem_total = torch.cuda.mem_get_info(device=device.torch_str()) mem_free, mem_total = torch.cuda.mem_get_info()
logger.debug( logger.debug(
"remaining CUDA VRAM usage: %s of %s", "remaining CUDA VRAM usage: %s of %s",
(mem_total - mem_free), (mem_total - mem_free),