Fix stack converter to support negative dims.
@@ -16,6 +16,10 @@ def convert_cat_trt7(ctx):
inputs = get_arg(ctx, 'input', pos=0, default=None)
dim = get_arg(ctx, 'dim', pos=1, default=0)
+ # Reverse negative dims.
+ if dim < 0:
+ dim = len(inputs[0].shape) - abs(dim + 1)
+
output = ctx.method_return
trt_inputs = [unsqueeze(ctx, i, dim) for i in inputs]
@@ -23,6 +27,7 @@ def convert_cat_trt7(ctx):
layer.axis = dim - 1
output._trt = layer.get_output(0)
+
class Stack(torch.nn.Module):
def __init__(self, dim):
super(Stack, self).__init__()
@@ -31,10 +36,27 @@ class Stack(torch.nn.Module):
def forward(self, *x):
return torch.stack(x, dim=self.dim)
+
@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 4, 4), (1, 4, 4)], enabled=trt_version() >= '7.0')
def test_Stack_basic_trt7():
return Stack(3)
+
@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 4, 4), (1, 4, 4)], enabled=trt_version() >= '7.0')
def test_Stack_basic2_trt7():
return Stack(1)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 4, 4), (1, 4, 4)], enabled=trt_version() >= '7.0')
+def test_Stack_neg1_dim_trt7():
+ return Stack(-1)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 4, 4), (1, 4, 4)], enabled=trt_version() >= '7.0')
+def test_Stack_neg2_dim_trt7():
+ return Stack(-2)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 4, 4), (1, 4, 4)], enabled=trt_version() >= '7.0')
+def test_Stack_neg3_dim_trt7():
+ return Stack(-3)
