Extract members from tensor while keeping the shape as it is

I have a torch tensor of size torch.Size([10, 1, 224, 224]), holding 10 images of size 224x224. I would like to extract single images but keep the tensor size as torch.Size([1, 1, 224, 224]). How can I achieve this?


To do this, I usually use .narrow() and extract a single element. For example to get the second image: data.narrow(0, 1, 1)

Hi @albanD! Thank you for the fast and correct response!