如何在PyTorch中为批量Tensor的指定索引位置赋值

这是批量Tensor操作里很常见的需求，咱们用PyTorch的高级索引就能优雅实现，完全不用写笨循环~

核心思路 #

要实现「每个批量元素对应修改自己的指定索引位置」，关键是要同时定位到批量维度的位置和每个子Tensor内的目标索引，然后批量赋值。

具体步骤&代码示例 #

假设我们的输入是：

• 批量Tensor t，形状为[N, D]（N是批量大小，D是每个子Tensor的维度）
• 索引Tensor indices，形状为[N]，每个元素是对应子Tensor里要修改的位置索引
• 要设置的目标值 X

直接看完整实现代码：

import torch

# 先准备示例数据
N = 3  # 批量大小
D = 5  # 每个子Tensor的维度
X = 99.0  # 要替换的目标值

# 随机生成原Tensor
t = torch.randn(N, D)
# 每个批量元素对应的要修改的索引
indices = torch.tensor([1, 3, 0])

# 1. 先克隆原Tensor，避免修改原数据（如果允许原地修改可以跳过这步）
t0 = t.clone()
# 2. 生成批量维度的索引：0到N-1，和indices一一对应
batch_indices = torch.arange(N, device=t.device)
# 3. 批量赋值：定位每个(batch_idx, idx)的位置，设置为X
t0[batch_indices, indices] = X

# 打印对比结果
print("原Tensor t:")
print(t)
print("\n修改后的Tensor t0:")
print(t0)

为什么这么做？ #

PyTorch的高级索引规则里，当你传入两个形状相同的一维Tensor（这里batch_indices和indices都是[N]），会自动把它们的元素一一配对，选取t0[0, indices[0]]、t0[1, indices[1]]……这些位置的元素，然后批量赋值为X，完美匹配你的需求。

注意事项 #

• 如果你的原Tensor在GPU上，一定要让batch_indices和它在同一个设备上，代码里的device=t.device就是干这个的，避免设备不匹配报错。
• 如果不需要保留原Tensor，可以直接对t操作，不用clone()，但这样原数据会被覆盖，操作前想清楚~

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