From e03b637f543dabd823dfb5f516aacc2857629fee Mon Sep 17 00:00:00 2001 From: Sean Sube Date: Sat, 25 Feb 2023 22:32:01 -0600 Subject: [PATCH] fix(api): only run CUDA GC on CUDA devices --- api/onnx_web/utils.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/api/onnx_web/utils.py b/api/onnx_web/utils.py index feff4f50..69012af9 100644 --- a/api/onnx_web/utils.py +++ b/api/onnx_web/utils.py @@ -89,7 +89,7 @@ def run_gc(devices: List[DeviceParams] = None): gc.collect() 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) with torch.cuda.device(device.torch_str()): torch.cuda.empty_cache()