Skip to content

Commit

Permalink
Fix a bug where in image_classification pipelines, a model_dir was no…
Browse files Browse the repository at this point in the history
…t being accepted properly (#1070)
  • Loading branch information
rahul-tuli authored Jun 15, 2023
1 parent 228751f commit 8257ca9
Showing 1 changed file with 5 additions and 0 deletions.
5 changes: 5 additions & 0 deletions src/deepsparse/image_classification/pipelines.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
Image classification pipeline
"""
import json
from pathlib import Path
from typing import Dict, List, Optional, Tuple, Type, Union

import numpy
Expand Down Expand Up @@ -135,6 +136,10 @@ class properties into an inference ready onnx file to be compiled by the
:return: file path to the ONNX file for the engine to compile
"""

model_path_: Path = Path(self.model_path)
if model_path_.is_dir():
return model_to_path(str(model_path_ / "model.onnx"))

return model_to_path(self.model_path)

def process_inputs(self, inputs: ImageClassificationInput) -> List[numpy.ndarray]:
Expand Down

0 comments on commit 8257ca9

Please sign in to comment.