pytorch : How to stack 2 tensors

Viewed 104

I’m trying to stack 2 tensors A.shape=(64,16,16) and B.shape=(64,16,16) in a tensor of shape C.shape=(1,128,16,16)

and non of functions I’ve tried work where
torch.stack => C.shape=(2,64,16,16) and torch.cat => C.shape=(128,16,16)

can enyones help me

1 Answers

Concat first and then use unsqueeze to add singleton dimension at 0th position

torch.cat([A, B]).unsqueeze(0)
Related