obtain matrix from a tensor in pytorch with arbitrary number of trailing dimensions

Question

I have a pytorch tensor with an arbitrary number of dimensions : ...X,Y,Z

I would like to have a function such that I give a number C, and I get the ...,C,Y,Z

my_matrix = [:,:,C,:,:]

But I dont know how many trailing dimensions are before C, I saw an answer with using tuples of slices but can seem to get it to work. Partial slices in pytorch / numpy with arbitrary and variable number of dimensions

Answer 1

I think ellipsis will do the job:

t = torch.randn(2, 3, 6, 5, 9, 3)
t[..., 4, :, :]

u = torch.randn(11, 4, 2, 7)
u[..., 2, :, :].shape