You can use torch.diag_embed 18 with torch.view(batch_size, -1) as the input. If your tensor is not necessarily contiguous, you can use torch.resize instead of torch.view.
Diagonals seems to be one of the bits of the numpy api that isn’t thought out terribly well w.r.t. being flexible/intuitive for use-cases like batching (see the default dimension behaviour for numpy.diagonal always seemed odd to me, too)…
Best regards
Thomas