[Python API] Added exception handling for 'None' model argument in save_model function (#21611)

This commit is contained in:
Areeb Syed authored and GitHub committed 2023-12-15 11:36:57 +01:00
1 parent 6fdd0418ad
commit 948b7405c7
2 files changed
+10

No files matched your search

@@ -165,6 +165,9 @@ PYBIND11_MODULE(_pyopenvino, m) {
[](std::shared_ptr<ov::Model>& model,
const py::object& xml_path,
bool compress_to_fp16) {
if (model == nullptr) {
throw py::attribute_error("'model' argument is required and cannot be None.");
}
ov::save_model(model,
Common::utils::convert_path_to_string(xml_path),
compress_to_fp16);
@@ -21,6 +21,7 @@ from openvino import (
set_batch,
get_batch,
serialize,
save_model,
)
from openvino.runtime import Output
from openvino.runtime.op.util import VariableInfo, Variable
@@ -566,3 +567,9 @@ def test_model_add_remove_variable():
assert variable_by_id.info.variable_id == "var_id_667"
model.remove_variable(variable_1)
assert len(model.get_variables()) == 1
def test_save_model_with_none():
with pytest.raises(AttributeError) as e:
save_model(model=None, output_model="model.xml")
assert "'model' argument is required and cannot be None." in str(e.value)