Hi, @zhl515.
Thanks for your solution and I’ve coded as you post, like this:
def kl_divergence(mu1, mu2, sigma_1, sigma_2):
	sigma_diag_1 = np.eye(sigma_1.shape[0]) * sigma_1
	sigma_diag_2 = np.eye(sigma_2.shape[0]) * sigma_2

	sigma_diag_2_inv = np.linalg.inv(sigma_diag_2)

	kl = 0.5 * (np.log(np.linalg.det(sigma_diag_2) / np.linalg.det(sigma_diag_2))
				- mu1.shape[0] + np.trace(np.matmul(sigma_diag_2_inv, sigma_diag_1))
				+ np.matmul(np.matmul(np.transpose(mu2 - mu1), sigma_diag_2_inv), (mu2 - mu1))
				)

	return kl

But I have another problem that I cannot call it in a “vectorized” way, or say, I don’t know how to let my code support batch operation in Pytorch. For instance, I have a tensor mu_1 with shape [batch_size,n] and another tensor mu_2 with shape [batch_size,n]. The variance tensors are same with mu. The n here is the dimensions of Gaussian distribution. How can I calculate kl loss for each instance without using for-loop ?
Any ideas? thx