add input_size & batch_size to onnx_export

2023-09-21 15:59:19 +09:00 · 2023-09-21 15:59:19 +09:00 · a9c29a05de
parent 245ad4f41a
commit a9c29a05de
1 changed files with 2 additions and 0 deletions
--- 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,
    )