如何对PyTorch Tensor的不同列按照自定义区间规则更新赋值？

start: end格式的切片仅支持传入单个整型/布尔值作为索引，无法直接传入一维张量实现按列差异化切片，因此触发类型错误。

可通过PyTorch广播机制生成布尔掩码完成批量赋值，无需手动遍历行列，全程调用内置操作，支持CPU/GPU端高效运行。
实现步骤如下：

1. 将定义的起止区间张量转换为PyTorch要求的64位整型索引格式
2. 生成行索引张量并调整维度适配广播规则
3. 按区间规则生成布尔掩码，对应需要赋值为1的位置
4. 基于掩码完成批量赋值

完整代码示例 #

import torch

# 初始化原始3行4列全2张量
n = torch.full((3, 4), 2.)
# 定义各列赋值的起止区间（左闭右开），统一为int64类型适配索引要求
start = torch.tensor([0, 0, 0, 0], dtype=torch.int64)
end = torch.tensor([2, 1, 3, 2], dtype=torch.int64)

# 生成行索引，维度调整为[行数, 1]以匹配列维度的广播规则
row_idx = torch.arange(n.shape[0]).unsqueeze(1)
# 生成布尔掩码：行索引落在对应列[start, end)区间内的位置标记为True
mask = (row_idx >= start) & (row_idx < end)
# 批量赋值
n[mask] = 1.

运行结果 #

输出张量符合预期：

tensor([[1., 1., 1., 1.],
        [1., 2., 1., 1.],
        [2., 2., 1., 2.]])

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