fix: Adjust reflection pad test cases to prevent runtime errors (#3088)

pytorch · Aug 16, 2024 · 0d7ca6f · 0d7ca6f
1 parent eb9f37f
commit 0d7ca6f
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 5 deletions.
diff --git a/py/torch_tensorrt/dynamo/conversion/impl/pad.py b/py/torch_tensorrt/dynamo/conversion/impl/pad.py
@@ -79,7 +79,7 @@ def get_padded_shape_tensors(
                     ).get_output(0)
                 )
         padded_shape_tensor = impl.cat.cat(
-            ctx, target, source_ir, f"{name}_cat", slices, 0
+            ctx, target, source_ir, f"{name}_cat_dim_{i}", slices, 0
         )
 
     start_indices_tensor = get_trt_tensor(

diff --git a/tests/py/dynamo/conversion/test_pad_aten.py b/tests/py/dynamo/conversion/test_pad_aten.py
@@ -146,9 +146,9 @@ def forward(self, input):
         [
             (
                 "4d",
-                (1, 1, 1, 1),
                 (2, 2, 2, 2),
-                (3, 3, 3, 3),
+                (4, 4, 4, 4),
+                (6, 6, 6, 6),
                 torch.float,
                 (1, 1, 2, 2),
             ),
@@ -199,9 +199,9 @@ def forward(self, input):
         [
             (
                 "5d",
-                (1, 1, 1, 1, 1),
                 (2, 2, 2, 2, 2),
-                (3, 3, 3, 3, 3),
+                (4, 4, 4, 4, 4),
+                (6, 6, 6, 6, 6),
                 torch.float,
                 (1, 2, 2, 1, 1, 2),
             ),