torch.cat vs. torch.stack

'torch.cat' and 'torch.stack' are both used to concatenate or stack tensors, but they work in slightly different ways.

'torch.cat' is used to concatenate tensors along a specific dimension. For example, if you have two tensors of shape (2, 3) and (2, 3), you can concatenate them along the first dimension (axis=0) to get a new tensor of shape (4, 3):

import torch
a = torch.randn(2, 3)
b = torch.randn(2, 3)
c = torch.cat((a, b), dim=0)
print(c.shape) # torch.Size([4, 3])

'torch.stack' is used to stack tensors along a new dimension. For example, if you have two tensors of shape (2, 3) and (2, 3), you can stack them along a new dimension (axis=0) to get a new tensor of shape (2, 2, 3):

import torch
a = torch.randn(2, 3)
b = torch.randn(2, 3)
c = torch.stack((a, b), dim=0)
print(c.shape) # torch.Size([2, 2, 3])

remember this, dimension argument in both torch.cat and torch.stack calls for a dimension to be fixed the called dimension will not change.