Shamoon
Shamoon

Reputation: 43491

How can I create a PyTorch tensor with all zeroes and a 1 in the middle of the third dimension?

I have a tensor, torch.Size([161, 161, 11]) and I want to set it all to zeroes, which I can do with: self.conv1.weight.data = torch.zeros(self.conv1.weight.data.size())

Except, I want column 6 (the middle) of the third dimension to be all ones. How do I do that?

Upvotes: 0

Views: 2384

Answers (1)

a_guest
a_guest

Reputation: 36239

You can assign it afterwards:

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

Or in case this tensor is trainable:

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

Upvotes: 2

Related Questions