Fix cat converter to support negative dims.
@@ -7,6 +7,10 @@ def convert_cat(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)
+
output = ctx.method_return
trt_inputs = add_missing_trt_tensors(ctx.network, inputs)
trt_inputs = broadcast_trt_tensors(ctx.network, trt_inputs, len(output.shape) - 1)
@@ -15,6 +19,7 @@ def convert_cat(ctx):
layer.axis = dim - 1
output._trt = layer.get_output(0)
+
class Cat(torch.nn.Module):
def __init__(self, dim):
super(Cat, self).__init__()
@@ -23,6 +28,17 @@ class Cat(torch.nn.Module):
def forward(self, *x):
return torch.cat(x, dim=self.dim)
+
@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 3, 4), (1, 17, 4)])
def test_Cat_basic():
return Cat(1)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 4, 4), (1, 4, 4)])
+def test_Cat_neg1_dim():
+ return Cat(-1)
+
+
+@add_module_test(torch.float32, torch.device('cuda'), [(1, 4, 4), (1, 4, 4), (1, 4, 4)])
+def test_Cat_neg2_dim():
+ return Cat(-2)
