What is the PyTorch equivalent of TensorFlow’s MultivariateNormalDiag distribution? Specifically, I have a B x N x D mean tensor and B x N x D variance tensor where B is batch size, N is number of data points, D is the dimension of each data point. I want to create a multi-variate normal distribution with diagonal covariance from these tensors. How can this be implemented ?