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])