Fixed ONNX size req padding (#2918)

chaiNNer-org · May 29, 2024 · 25e9dc3 · 25e9dc3
1 parent 7311cde
commit 25e9dc3
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/backend/src/nodes/impl/pytorch/convert_to_onnx_impl.py b/backend/src/nodes/impl/pytorch/convert_to_onnx_impl.py
@@ -24,8 +24,8 @@ def convert_to_onnx_impl(
         input_name: {0: "batch_size", 2: "height", 3: "width"},
         output_name: {0: "batch_size", 2: "height", 3: "width"},
     }
-    size = max(model.size_requirements.minimum, 3)
-    size = size + (size % model.size_requirements.multiple_of)
+    size = 3
+    size += model.size_requirements.get_padding(size, size)[0]
     dummy_input = torch.rand(1, model.input_channels, size, size)
     dummy_input = dummy_input.to(device)