SparseZoo V2 Compatibility (#1233)

neuralmagic · Sep 12, 2023 · da6ab6f · da6ab6f
1 parent d92822f
commit da6ab6f
Showing 1 changed file with 6 additions and 5 deletions.
diff --git a/tests/deepsparse/image_classification/test_pipelines.py b/tests/deepsparse/image_classification/test_pipelines.py
@@ -59,12 +59,13 @@ def test_image_classification_pipeline_preprocessing(
     ic_pipeline = Pipeline.create("image_classification", model_path=zoo_stub)
     zoo_model = Model(zoo_stub)
     data_originals_path = None
-    if zoo_model.sample_originals is not None:
-        if not zoo_model.sample_originals.files:
-            zoo_model.sample_originals.unzip()
-        data_originals_path = zoo_model.sample_originals.path
+    if zoo_model.sample_inputs is not None:
+        if not zoo_model.sample_inputs.files:
+            zoo_model.sample_inputs.unzip()
+        data_originals_path = zoo_model.sample_inputs.path
     for idx, sample in enumerate(load_numpy_list(data_originals_path)):
-        image_raw = list(sample.values())[0]
+        image_raw = list(sample.values())[0].astype(numpy.uint8)
+        image_raw = numpy.transpose(image_raw, axes=[1, 2, 0])
         image_raw = Image.fromarray(image_raw)
 
         preprocessed_image_pipeline = ic_pipeline.process_inputs(