spamosaic.loss.set_diag

spamosaic.loss.set_diag(matrix, v)

Set the diagonal of a square tensor to a specified value.

Parameters:

• matrix (torch.Tensor) – 2D square tensor of shape [N, N] whose diagonal will be modified.

• v (float) – Value to set on the diagonal.

Returns:

The same tensor with its diagonal set to v.

Return type:

torch.Tensor