我有一个6维的全零PyTorch张量lrel_w，我想在前三个维度和后三个维度的索引相匹配的位置填充1s。我目前使用3个嵌套的for循环简单地解决了这个问题：

lrel_w = torch.zeros(  input_size[0], input_size[1], input_size[2],  input_size[0], input_size[1], input_size[2])for c in range(input_size[0]):  for x in range(input_size[1]):    for y in range(input_size[2]):      lrel_w[c,x,y,c,x,y] = 1

我相信一定有更有效的方法来做这件事，但我还没有找到解决方案。

回答：

你可以试试这个方法。

import torchc, m, n = input_size[0], input_size[1], input_size[2]t = torch.zeros(c, m, n, c, m, n)i, j, k = torch.meshgrid(torch.arange(c), torch.arange(m), torch.arange(n))i = i.flatten()j = j.flatten()k = k.flatten()t[i, j, k, i, j, k] = 1

如果你需要参考，以下是meshgrid的工作原理说明。