Remove batch dimension from samples (#1673)

neuralmagic · Jul 17, 2023 · e7f78e0 · e7f78e0
1 parent 7333f6b
commit e7f78e0
Showing 1 changed file with 3 additions and 2 deletions.
diff --git a/src/sparseml/yolov8/utils/export_samples.py b/src/sparseml/yolov8/utils/export_samples.py
@@ -163,15 +163,16 @@ def _export_ort_outputs(
     # Run model to get onnxruntime outputs
     ort_inputs = {session.get_inputs()[0].name: image}
     ort_outs = session.run(None, ort_inputs)
-    preds = ort_outs
+    preds = numpy.squeeze(ort_outs, axis=0)
+    preds = numpy.squeeze(preds, axis=0)
 
     sample_output_filename = os.path.join(sample_out_dir, f"out-{file_idx}.npz")
     numpy.savez(sample_output_filename, preds)
 
 
 def _export_inputs(image: torch.Tensor, sample_in_dir: str, file_idx: str):
 
-    sample_in = image.detach().to("cpu")
+    sample_in = image.detach().to("cpu").squeeze(0)
 
     sample_input_filename = os.path.join(sample_in_dir, f"inp-{file_idx}.npz")
     numpy.savez(sample_input_filename, sample_in)