Stack a vector and another tensor of size 0 or 2

1

I have a which is a tensor object with shape (3, ).
I have b which is a tensor object with shape either [] (scalar) or (2, ).

I want to concatenate a and b into a single vector of length (4, ) or (5, ).
It seems that torch.cat or torch.stack will fail doing so.

Is there other one liner option to do without if to handle the case of the scalar?

2

I found something to work:

import torch

a = torch.rand(3)
b1 = torch.rand(2)
b2 = b1[0]

torch.cat((a, torch.atleast_1d(b1)), dim = 0)
torch.cat((a, torch.atleast_1d(b2)), dim = 0)

If there a better way, I am all ears.

1 Like