How to keep the tensor size the same?

1 
A = torch.rand(4,3,5,5)

B = A[:,0,:,:]

B.size()
Out[33]: torch.Size([4, 5, 5])

is there anyway other than using torch.unsqueeze to keep the size of tensor B to be [4,1,3,3]?

2

Capture

3

hahah
I was more interested in some keepdim = true way, but that works too i guess
Thanks