我如何创建一个PyTorch张量全为零，在第三维的中间为1？

Question

我有一个张量torch.Size([161, 161, 11])，我想将其全部设置为零，我可以使用：self.conv1.weight.data = torch.zeros(self.conv1.weight.data.size())

除了，我希望第三维的第6列（中间）全为1。我该怎么办？

Answer 1

您可以在以后分配它：

self.conv1.weight.data[:, :, 6] = 1.0

或者该张量是可训练的：

with torch.no_grad():
    self.conv1.weight.data[:, :, 6] = 1.0