How to compute a 2-D square tensor with the elements of input as the diagonal

Hi all,

We have a problem when using the API torch.diag(input, diagonal=0, out=None).

If the input is a-2D tensor, and input.size(0) = batch_size. How can we compute the square tensor
with the elements of input as as the diagonal.

Thanks.