fix(api): correct device for VRAM check

ssube · Feb 17, 2023 · cf8ca34 · cf8ca34
1 parent 56a8262
commit cf8ca34
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py
@@ -144,7 +144,7 @@ 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())
+                mem_free, mem_total = torch.cuda.mem_get_info()
                 logger.debug(
                     "remaining CUDA VRAM usage: %s of %s",
                     (mem_total - mem_free),