title: How to construct a 3D Tensor where every 2D sub tensor is a diagonal matrix in PyTorch?
body: <p>Consider I have 2D Tensor, <code>index_in_batch * diag_ele</code>.
  How can I get a 3D Tensor <code>index_in_batch * Matrix</code> (who is a diagonal matrix, construct by drag_ele)?</p>
  
  <p>The <code>torch.diag()</code> construct diagonal matrix only when input is 1D, and return diagonal element when input is 2D.</p>