I believe you are looking for the ‘dim’ argument of torch.cat



dim (int, optional) – the dimension over which the tensors are concatenated


>>> a = torch.rand((124,39,1024))
>>> b = torch.rand((124,39,1024))
>>> torch.cat((a,b)).shape
torch.Size([248, 39, 1024])

>>> torch.cat((a,b), dim=0).shape
torch.Size([248, 39, 1024])

>>> torch.cat((a,b), dim=1).shape
torch.Size([124, 78, 1024])

>>> torch.cat((a,b), dim=2).shape
torch.Size([124, 39, 2048])