torch.stack and torch.cat both ‘concatenate’ Tensors but the key different is that torch.stack applies this concatenation along a new dimenion, thereby stacking the Tensors whereas torch.cat does a direct concatenation along already existing dimensions
For example,
a = torch.randn(7,768)
b = torch.randn(8,768)
torch.cat((a,b)).shape #returns torch.Size([15, 768])
Trying torch.stack results in the error,
RuntimeError: stack expects each tensor to be equal size, but got [7, 768] at entry 0 and [8, 768] at entry 1
because torch.stack takes Tensors of the same shape and stacks them into a larger Tensor. You can think of having a B number of matrices of size [N, N] and stacking them together to get a Tensor [B,N,N].
I’ll repeat your original problem with two new tensors that should explain the difference between torch.stack and torch.cat.
For example,
aa = torch.randn(4,768)
bb = torch.randn(4,768)
torch.stack((aa,bb)).shape #returns torch.Size([2, 4, 768]) #stacked along new dimension
torch.cat((aa,bb)).shape #returns torch.Size([8, 768]) #concatenated along existing dimension
This example should clearly show the difference between torch.stack and torch.cat. 