spamosaic.loss.set_diag

spamosaic.loss.set_diag(matrix, v)[source]

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

Parameters:

matrix (torch.Tensor) – A 2D square tensor whose diagonal will be modified.
v (float) – Value to set on the diagonal.

Returns:

The modified tensor with updated diagonal values.

Return type:

torch.Tensor