apextoaster
/
onnx-web
1
0
Fork 0
Code Actions

fix(api): only run CUDA GC on CUDA devices

This commit is contained in:
Sean Sube 2023-02-25 22:32:01 -06:00
parent 6809d2da82
commit e03b637f54
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

@ -89,7 +89,7 @@ def run_gc(devices: List[DeviceParams] = None):
gc.collect() gc.collect()
if torch.cuda.is_available() and devices is not None: if torch.cuda.is_available() and devices is not None:
for device in devices: for device in [d for d in devices if d.device.startswith("cuda")]:
logger.debug("running Torch garbage collection for device: %s", device) logger.debug("running Torch garbage collection for device: %s", device)
with torch.cuda.device(device.torch_str()): with torch.cuda.device(device.torch_str()):
torch.cuda.empty_cache() torch.cuda.empty_cache()