How to obtain a 1D output from a NN using 2D input (batches)?

Hi all,

I have the following setting: I want to learn a gaussian distribution via two NNs, i.e., one NN maps onto the mean vector and the other NN maps onto the covariance matrix. My issue is that I am feeding batches of shape (batch_dim, observation_dim) into the NN and hence obtain for the mean vector also a vector of size (batch_dim, observation_dim) and for the covariance matrix a tensor of shape (batch_dim, observation_dim * observation_dim). However, I do not want to have an individual mean vector and covariance matrix for each sample in the batch, but instead one mean vector and one covariance matrix for the whole batch. I.e., I want to get rid of the batch_dim in the output: The mean vector should have shape (observation_dim) and the covariance matrix should have shape (observation_dim, observation_dim). How can this be achieved?

Thank you in advance!