From a9c29a05de651eb700daef24a47437e02770c496 Mon Sep 17 00:00:00 2001
From: Youssef Boulaoaune <43298428+Boulaouaney@users.noreply.github.com>
Date: Thu, 21 Sep 2023 15:59:19 +0900
Subject: [PATCH] add input_size & batch_size to onnx_export

---
 onnx_export.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/onnx_export.py b/onnx_export.py
index 3baab369ca..554d314b0b 100644
--- 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,
     )