Suppose I have a tensor of some arbitrary, unknown length. Take for instance:
a = torch.randn(A, B, D)

or
b = torch.randn(A, B, C, D)

where A to D are arbitrary.
How do I index the tensor such that I get the last element of X,  which is always the second to last index, producing two tensors of the remaining elements?
a_x = foo(a)
b_x = foo(b)

torch.allclose(a[:,:,X-1,:], a_x)
> True

torch.allclose(b[:,:,:,X-1,:], b_x)
> True

a_x.shape
> torch.Size([A,B,D])

b_x.shape
> torch.Size([A,B,C,D])