"torch.diagonal()" has different default dimensions, so those need
   to be explicitly specified.
Parameters:
      * input (Tensor) -- the input tensor. Must be at least
        1-dimensional.
  * **offset** (*int**, **optional*) -- which diagonal to
    consider. Default: 0 (main diagonal).

  * **dim1** (*int**, **optional*) -- first dimension with respect
    to which to take diagonal. Default: -2.

  * **dim2** (*int**, **optional*) -- second dimension with
    respect to which to take diagonal. Default: -1.

Example:
  >>> a = torch.randn(2, 3)
  >>> torch.diag_embed(a)
  tensor([[[ 1.5410,  0.0000,  0.0000],
           [ 0.0000, -0.2934,  0.0000],
           [ 0.0000,  0.0000, -2.1788]],

          [[ 0.5684,  0.0000,  0.0000],
           [ 0.0000, -1.0845,  0.0000],
           [ 0.0000,  0.0000, -1.3986]]])

  >>> torch.diag_embed(a, offset=1, dim1=0, dim2=2)
  tensor([[[ 0.0000,  1.5410,  0.0000,  0.0000],
