如何高效移除Tensor中的特定值并规整维度？以值1为例

给定原始Tensor：

tensor([[2, 5, 4, 2, 1, 1, 1, 1, 5],
        [4, 0, 1, 5, 0, 3, 5, 1, 4],
        [1, 1, 7, 2, 1, 5, 1, 1, 1],
        [5, 5, 5, 3, 0, 4, 3, 1, 1],
        [1, 1, 1, 1, 5, 3, 2, 7, 3],
        [6, 6, 1, 1, 1, 1, 1, 1, 1]])

需要尽可能移除其中的1，转换为如下目标Tensor：

tensor([[2., 5., 4., 2., 5., 1., 1.],
        [4., 0., 5., 0., 3., 5., 4.],
        [7., 2., 5., 1., 1., 1., 1.],
        [5., 5., 5., 3., 0., 4., 3.],
        [5., 3., 2., 7., 3., 1., 1.],
        [6., 6., 1., 1., 1., 1., 1.]])

你当前的逻辑是可行的：先统计每行非1元素的数量，取最大值初始化全1的目标Tensor，再逐行填充非1元素。但逐行循环在数据量大的时候效率偏低，这里提供一种无循环的向量化实现，利用PyTorch的广播和索引操作提速。

import torch

t = torch.tensor([[2, 5, 4, 2, 1, 1, 1, 1, 5],
                  [4, 0, 1, 5, 0, 3, 5, 1, 4],
                  [1, 1, 7, 2, 1, 5, 1, 1, 1],
                  [5, 5, 5, 3, 0, 4, 3, 1, 1],
                  [1, 1, 1, 1, 5, 3, 2, 7, 3],
                  [6, 6, 1, 1, 1, 1, 1, 1, 1]])

# 标记所有非1元素的位置
mask = t != 1
# 统计每行非1元素的个数
non_one_count = mask.sum(dim=1)
max_len = non_one_count.max()

# 生成目标Tensor的行列索引矩阵
row_indices = torch.arange(t.size(0)).unsqueeze(1).repeat(1, max_len)
col_indices = torch.arange(max_len).unsqueeze(0).repeat(t.size(0), 1)

# 标记哪些位置需要填充非1元素
valid_mask = col_indices < non_one_count.unsqueeze(1)

# 初始化全1的目标Tensor
desired_tensor = torch.ones(t.size(0), max_len, dtype=t.dtype).float()
# 收集所有非1元素
non_one_elements = t[mask]
# 一次性填充所有有效位置
desired_tensor[valid_mask] = non_one_elements

print(desired_tensor)

• 用mask快速定位所有非1元素，避免逐行遍历；
• valid_mask生成和目标Tensor同形状的布尔矩阵，明确需要填充非1元素的位置；
• 利用PyTorch的高级索引一次性完成填充，完全取消循环，在数据规模较大时效率提升显著。

内容的提问来源于stack exchange，提问作者amy