KL-divergence between two multivariate gaussian

[NILOIDE (Nil Stolt Ansó)]

Check this thread.
torch.diag_embed() should do the trick.