diff --git a/machine-learning/app/models/base.py b/machine-learning/app/models/base.py index 496012270..16adc159a 100644 --- a/machine-learning/app/models/base.py +++ b/machine-learning/app/models/base.py @@ -186,7 +186,7 @@ class InferenceModel(ABC): case "CPUExecutionProvider" | "CUDAExecutionProvider": option = {"arena_extend_strategy": "kSameAsRequested"} case "OpenVINOExecutionProvider": - option = {"device_type": "GPU_FP32"} + option = {"device_type": "GPU_FP32", "cache_dir": (self.cache_dir / "openvino").as_posix()} case _: option = {} options.append(option) diff --git a/machine-learning/app/test_main.py b/machine-learning/app/test_main.py index d79da0cbd..22038a654 100644 --- a/machine-learning/app/test_main.py +++ b/machine-learning/app/test_main.py @@ -88,7 +88,7 @@ class TestBase: encoder = OpenCLIPEncoder("ViT-B-32__openai", providers=["OpenVINOExecutionProvider", "CPUExecutionProvider"]) assert encoder.provider_options == [ - {"device_type": "GPU_FP32"}, + {"device_type": "GPU_FP32", "cache_dir": (encoder.cache_dir / "openvino").as_posix()}, {"arena_extend_strategy": "kSameAsRequested"}, ]