Skip to content

Commit

Permalink
Update model_importer.py
Browse files Browse the repository at this point in the history
  • Loading branch information
gkrivor authored Mar 22, 2024
1 parent 0fd85e7 commit 29ad1d9
Showing 1 changed file with 11 additions and 0 deletions.
11 changes: 11 additions & 0 deletions src/frontends/onnx/tests/tests_python/utils/model_importer.py
Original file line number Diff line number Diff line change
Expand Up @@ -141,16 +141,27 @@ def _execute_pb_data(
executed_tests = executed_tests + 1
return executed_tests


def _add_model_import_test(self, model_test: ExtOnnxTestCase) -> None:
# model is loaded at runtime, note sometimes it could even
# never loaded if the test skipped
model_marker = [None] # type: List[Optional[Union[ModelProto, NodeProto]]]

def run_import(test_self: Any, device: Text) -> None:
model = ModelImportRunner._load_onnx_model(model_test.model_dir, model_test.model)
model_marker[0] = model_test.model_dir / model_test.model
assert import_onnx_model(model)

self._add_test("ModelImport", model_test.name, run_import, model_marker)

def _add_model_execution_test(self, model_test: ExtOnnxTestCase) -> None:
# model is loaded at runtime, note sometimes it could even
# never loaded if the test skipped
model_marker = [None] # type: List[Optional[Union[ModelProto, NodeProto]]]

def run_execution(test_self: Any, device: Text) -> None:
model = ModelImportRunner._load_onnx_model(model_test.model_dir, model_test.model)
model_marker[0] = model_test.model_dir / model_test.model
prepared_model = self.backend.prepare(model, device)
assert prepared_model is not None
executed_tests = ModelImportRunner._execute_npz_data(
Expand Down

0 comments on commit 29ad1d9

Please sign in to comment.