如何用PyTorch函数修改2D张量每行指定索引后的所有元素？

当然可以不用for循环，靠PyTorch的广播和掩码机制就能高效完成这个需求。核心思路是生成一个掩码矩阵，标记出每行需要保留的元素位置，再批量修改不需要保留的部分。

实现代码 #

import torch

def do_something(A, idx, fill_value):
    # 生成与A同形状的列索引矩阵
    col_indices = torch.arange(A.size(1), device=A.device).expand_as(A)
    # 生成掩码：每行中列索引 <= 对应行idx值的位置保留原数据
    mask = col_indices <= idx.unsqueeze(1)
    # 克隆原张量后，将掩码外的位置设为指定填充值
    B = A.clone()
    B[~mask] = fill_value
    return B

测试验证 #

用你给出的示例数据测试：

A = torch.tensor([[ 0.0265,  0.1408,  1.6343,  0.3062,  0.3950],
        [ 0.6614,  1.6490, -1.8821, -0.2621,  1.4919],
        [ 1.0964,  1.5243,  0.5565, -0.7764, -1.2834],
        [-1.2818, -0.9590, -0.8375, -1.3218,  0.1081],
        [ 1.4578,  0.6749, -1.7441,  0.6515,  0.7278]])
idx = torch.tensor([0, 1, 2, 1, 3])
B = do_something(A, idx, 0)
print(B)

输出结果和示例完全一致：

tensor([[ 0.0265,  0.0000,  0.0000,  0.0000,  0.0000],
        [ 0.6614,  1.6490,  0.0000,  0.0000,  0.0000],
        [ 1.0964,  1.5243,  0.5565,  0.0000,  0.0000],
        [-1.2818, -0.9590,  0.0000,  0.0000,  0.0000],
        [ 1.4578,  0.6749, -1.7441,  0.6515,  0.0000]])

原理说明 #

1. 列索引矩阵：用torch.arange生成0到列数-1的序列，再通过expand_as扩展成和A一样的形状，这样每行都是完整的列索引。
2. 掩码生成：把1D的idx通过unsqueeze(1)转成列向量，利用PyTorch的广播机制和列索引矩阵做比较，得到每行中哪些位置需要保留（列索引≤对应行的idx值）。
3. 批量修改：克隆原张量后，用取反的掩码~mask定位到需要修改的位置，直接赋值为填充值，全程无循环，效率远高于逐行遍历。

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