我有两个 torch 张量
mask = torch.ones(1024, 64, dtype=torch.float32)
indices = torch.randint(0, 64, (1024, ))
对于mask中的每i行，我希望将indices的第i个元素指定的索引之后的所有元素设置为零.例如，如果indices的第一个元素是50，那么我想设置mask[0, 50:]=0.可以在不使用for循环的情况下实现这一点吗？
使用for循环的解决方案:
for i in range(mask.shape[0]):
mask[i, indices[i]:] = 0
您可以首先生成一个大小为(1024x64)的张量，其中每行的数字排列在0到63之间.然后使用reshape 为(1024x1)的索引应用逻辑运算
mask = torch.ones(1024, 64, dtype=torch.float32)
indices = torch.randint(0, 64, (1024, 1)) # Note the dimensions
mask[torch.arange(0, 64, dtype=torch.float32).repeat(1024,1) >= indices] = 0