revert

logivations · Aug 6, 2024 · 6de1172 · 6de1172
1 parent 2de5f17
commit 6de1172
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/auto_training/export.py b/auto_training/export.py
@@ -129,11 +129,11 @@ def parse_args():
 
     if args.shape is None:
         img_scale = cfg.test_pipeline[1]['img_scale']
-        input_shape = (4, 3, img_scale[1], img_scale[0])
+        input_shape = (1, 3, img_scale[1], img_scale[0])
     elif len(args.shape) == 1:
-        input_shape = (4, 3, args.shape[0], args.shape[0])
+        input_shape = (1, 3, args.shape[0], args.shape[0])
     elif len(args.shape) == 2:
-        input_shape = (4, 3) + tuple(args.shape)
+        input_shape = (1, 3) + tuple(args.shape)
     else:
         raise ValueError('invalid input shape')