torch.masked_select¶
- torch.masked_select(input, mask, *, out=None) Tensor ¶
返回一个新的 1-D 张量,该张量根据布尔掩码
mask
(它是一个 BoolTensor)索引input
张量。mask
张量和input
张量的形状不需要匹配,但它们必须是 可广播的。注意
返回的张量不使用与原始张量相同的存储空间
示例
>>> 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])