[FIXED] How to add a new dim to a a pytorch tensor?

Issue

In numpy I would do

a = np.zeros((4, 5, 6))
a = a[:, :, np.newaxis, :]
assert a.shape == (4, 5, 1, 6)

How to do the same in torch?

Solution

a = torch.zeros(4, 5, 6)
a = a[:, :, None, :]
assert a.shape == (4, 5, 1, 6)

Answered By - Gulzar