sevenhill/Immich
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

fix(ml): serialize trt-rtx sessions (#31963)

Browse Source 
fix(ml): serialize trt-rtx sessions (#31951)

(cherry picked from commit a3c8b359f2fdc47699f1530215a23fe8bd34c454)

Co-authored-by: Mert <101130780+mertalev@users.noreply.github.com>
This commit is contained in:
immich-push-o-matic[bot]
2026-09-30 23:08:58 +00:00
committed by GitHub
co-authored by Mert
parent eace8c2809
commit a506c8a28f
2 changed files with 36 additions and 11 deletions
Download Patch File Download Diff File Expand all files Collapse all files
machine-learning/test_main.py
+8 -1
View File
@@ -295,6 +295,7 @@ class TestOrtSessions:
TRT_EP = ["TensorrtExecutionProvider", "CUDAExecutionProvider", "CPUExecutionProvider"]
ROCM_EP = ["MIGraphXExecutionProvider", "CPUExecutionProvider"]
COREML_EP = ["CoreMLExecutionProvider", "CPUExecutionProvider"]
TRT_RTX_EP = ["nv_tensorrt_rtx", "CUDAExecutionProvider", "CPUExecutionProvider"]
@pytest.mark.providers(CPU_EP)
def test_sets_cpu_provider(self, ort_session: mock.Mock, providers: list[str]) -> None:
@@ -341,6 +342,12 @@ class TestOrtSessions:
assert given_providers(ort_session) == self.COREML_EP
@pytest.mark.providers(TRT_RTX_EP)
def test_leaves_cuda_out_beside_tensorrt_rtx(self, ort_session: mock.Mock, providers: list[str]) -> None:
ort_sessions("ViT-B-32__openai")
assert given_providers(ort_session) == ["nv_tensorrt_rtx", "CPUExecutionProvider"]
def test_leaves_a_dimension_free_when_the_model_feeds_several_sizes(
self, ort_session: mock.Mock, mocker: MockerFixture
) -> None:
@@ -408,7 +415,7 @@ class TestOrtSessions:
events.append("open")
return mock.DEFAULT
def run(output_names: Any, feed: dict[str, np.ndarray]) -> list[np.ndarray]:
def run(output_names: Any, feed: dict[str, np.ndarray], run_options: Any = None) -> list[np.ndarray]:
events.append(f"run at {feed['image'].shape[2]}") # a dim the graph leaves free takes the shape's size
return [np.zeros(1)]