helpy4/immich
1
0
Fork 0
mirror of https://github.com/immich-app/immich synced 2025-10-17 18:19:27 +00:00
Code Issues Projects Releases Packages Wiki Activity Actions

feat(ml): round-robin device assignment (#13237)

Browse source 
* round-robin device assignment

* docs and tests

clarify doc
This commit is contained in:
Mert 2024-10-07 17:37:45 -04:00 • committed by GitHub
parent 063969ca05
commit bd826b0b9b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
8 changed files with 62 additions and 7 deletions
Download patch file Download diff file Expand all files Collapse all files

6
machine-learning/app/sessions/ort.py
View file

@ -86,11 +86,13 @@ class OrtSession:
provider_options = []
for provider in self.providers:
match provider:
case "CPUExecutionProvider" | "CUDAExecutionProvider":
case "CPUExecutionProvider":
options = {"arena_extend_strategy": "kSameAsRequested"}
case "CUDAExecutionProvider":
options = {"arena_extend_strategy": "kSameAsRequested", "device_id": settings.device_id}
case "OpenVINOExecutionProvider":
options = {
"device_type": "GPU",
"device_type": f"GPU.{settings.device_id}",
"precision": "FP32",
"cache_dir": (self.model_path.parent / "openvino").as_posix(),
}