torch.masked_select¶
- torch.masked_select(input, mask, *, out=None) Tensor ¶
Returns a new 1-D tensor which indexes the
input
tensor according to the boolean maskmask
which is a BoolTensor.The shapes of the
mask
tensor and theinput
tensor don’t need to match, but they must be broadcastable.Note
The returned tensor does not use the same storage as the original tensor
- Parameters:
input (Tensor) – the input tensor.
mask (BoolTensor) – the tensor containing the binary mask to index with
- Keyword Arguments:
out (Tensor, optional) – the output tensor.
Example:
>>> x = torch.randn(3, 4) >>> x tensor([[ 0.3552, -2.3825, -0.8297, 0.3477], [-1.2035, 1.2252, 0.5002, 0.6248], [ 0.1307, -2.0608, 0.1244, 2.0139]]) >>> mask = x.ge(0.5) >>> mask tensor([[False, False, False, False], [False, True, True, True], [False, False, False, True]]) >>> torch.masked_select(x, mask) tensor([ 1.2252, 0.5002, 0.6248, 2.0139])