Update infer_onnx.py

infer_onnx.py

@@ -25,6 +25,7 @@ import torchvision.transforms as transforms
 parser = argparse.ArgumentParser()
 parser.add_argument('--onnx_path', type=str, default="EfficientNet_int.onnx", required=False)
 parser.add_argument('--image_path', type=str, required=True)
+parser.add_argument('--data_format', type=str, choices=["nchw", "nhwc"], default="nchw")
 parser.add_argument(
     "--ipu",
     action="store_true",
@@ -51,6 +52,8 @@ def read_image():
     normalize,
   ])
   img_tensor = transform(image).unsqueeze(0)
+  if args.data_format == "nhwc":
+    img_tensor = transform(image).unsqueeze(0).transpose(1, 3)
   return img_tensor.numpy()