Height and width argument correction for random tensor creation

sovit-123 · May 25, 2024 · 8221971 · 8221971
1 parent 1ea102f
commit 8221971
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/export.py b/export.py
@@ -81,7 +81,7 @@ def main(args):
     model.eval()
 
     # Input to the model
-    x = torch.randn(1, 3, args['width'], args['height'])
+    x = torch.randn(1, 3, args['height'], args['width'])
 
     # Export the model
     torch.onnx.export(