Solved by iacob in post #4 
              
              
                You can do this with a for loop over the sub-tensors: 
# across dim0
for i in range(data.size(0)):
    data[i].fill_diagonal_(0)

If you need to perform this over an arbitrary two dimensions of a 3d tensor, simply apply the fill to the appropriate slices: 
# across dim1
for i in range(data.size(1)):…