Normally, you can slice a 2D tensor like this slice = t[:, :k] where k is an integer. Is it possible to do something like this but with k being a 1-dimensional vector of integers with the number of items that I want to obtain for each row?

Masking the items with 0’s or NaN would also be fine.

For example:

k = torch.Tensor([1,3,2])
t = torch.Tensor([1,1,1], [2,2,2], [3,3,3])
# perform some operations and the result should be
# 1 - -
# 2 2 2
# 3 3 -

k = torch.Tensor([1,3,2])
t = torch.Tensor([[1,1,1], [2,2,2], [3,3,3]])
h, w = t.shape
# create a mask of ones with size h, w+1
mask = torch.ones(h, w+1)
# set other elements to 0, depending on length of each row
mask[torch.arange(h), k.long()] = 0.
mask = mask.cumprod(dim=1)
# multiply the mask
result = t * mask[:, :-1]
print(result)