Создать маску среза в pytorch?

Question

Есть ли способ указать маску на основе операции среза?
Например

A = torch.arange(6).view((2,3))
# A = [[0,1,2], [3,4,5]]

mask_slice = torch.mask_slice(A[:,1:])
# mask_slice = [[0,1,1],[0,1,1]]

Answer 1

Вы можете сделать что-то вроде этого (если я правильно понял ваш вопрос):

mask_slice = torch.zeros(A.shape, dtype=bool)
mask_slice[:, 1:] = 1

# tensor([[False,  True,  True],
#         [False,  True,  True]])