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

ssube · Feb 26, 2023 · e03b637 · e03b637
1 parent 6809d2d
commit e03b637
Showing 1 changed file with 1 addition and 1 deletion.
diff --git 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()