Update infer_onnx.py

infer_onnx.py

@@ -36,7 +36,7 @@ parser.add_argument(
     default="vaip_config.json",
     help="Path of the config file for seting provider_options.",
 )
-
+parser.add_argument('--data_format', type=str, choices=["nchw", "nhwc"], default="nchw")
 args = parser.parse_args()
 
 
@@ -51,6 +51,8 @@ def read_image():
     normalize,
   ])
   img_tensor = transform(image).unsqueeze(0)
+  if args.data_format == "nhwc":
+    img_tensor = transform(image).unsqueeze(0).permute((0, 2, 3, 1))
   return img_tensor.numpy()