cjones/immich
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

add test

Browse Source
This commit is contained in:
mertalev
2025-04-19 16:30:50 -04:00
parent 543bc72ae3
commit 5a3b11d603
1 changed files with 7 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
machine-learning/test_main.py
+7
View File
@@ -180,6 +180,7 @@ class TestOrtSession:
CUDA_EP_OUT_OF_ORDER = ["CPUExecutionProvider", "CUDAExecutionProvider"] CUDA_EP_OUT_OF_ORDER = ["CPUExecutionProvider", "CUDAExecutionProvider"]
TRT_EP = ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"] TRT_EP = ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
ROCM_EP = ["ROCMExecutionProvider", "CPUExecutionProvider"] ROCM_EP = ["ROCMExecutionProvider", "CPUExecutionProvider"]
COREML_EP = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
@pytest.mark.providers(CPU_EP) @pytest.mark.providers(CPU_EP)
def test_sets_cpu_provider(self, providers: list[str]) -> None: def test_sets_cpu_provider(self, providers: list[str]) -> None:
@@ -225,6 +226,12 @@ class TestOrtSession:
assert session.providers == self.ROCM_EP assert session.providers == self.ROCM_EP
@pytest.mark.providers(COREML_EP)
def test_uses_coreml(self, providers: list[str]) -> None:
session = OrtSession("ViT-B-32__openai")
assert session.providers == self.COREML_EP
def test_sets_provider_kwarg(self) -> None: def test_sets_provider_kwarg(self) -> None:
providers = ["CUDAExecutionProvider"] providers = ["CUDAExecutionProvider"]
session = OrtSession("ViT-B-32__openai", providers=providers) session = OrtSession("ViT-B-32__openai", providers=providers)