add input_size & batch_size to onnx_export

huggingface · Sep 21, 2023 · a9c29a0 · a9c29a0
1 parent 245ad4f
commit a9c29a0
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/onnx_export.py b/onnx_export.py
@@ -90,6 +90,8 @@ def main():
         check_forward=args.check_forward,
         training=args.training,
         verbose=args.verbose,
+        input_size=(3, args.img_size, args.img_size),
+        batch_size=args.batch_size,
     )