Hello Kim,
how about the following?
x = torch.randn(2,5,5,1)
idx = torch.arange(0,5, out=torch.LongTensor())
x[:,idx,idx] = 0
print(x[0,:,:,0],x[1,:,:,0])

Best regards
Thomas
P.S.: I’m not sure about the performance characteristics of this vs. Peter’s solution.