Masking diagonal to a specific value with PyTorch tensors
How do I fill the diagonal with a value in torch? In numpy you can do:
a = np.zeros((3, 3), int)
np.fill_diagonal(a, 5)

array([[5, 0, 0],
       [0, 5, 0],
       [0, 0, 5]])

I know that torch.diag() returns the diagonal, but how to use this as a mask to assign new values is beyond me. I haven't been able to find the answer here or in the PyTorch documentation.

You can do this in PyTorch using fill_diagonal_:
&gt;&gt;&gt; a = torch.zeros(3, 3)
&gt;&gt;&gt; a.fill_diagonal_(5)
tensor([[5, 0, 0],
        [0, 5, 0],
        [0, 0, 5]])

