Regression tasks: Epistemic and Aleatoric Uncertainty Estimation

Hello, I am trying to use this repository for regression tasks (I can see the examples seem to focus on classification tasks).

I would like to do estimate epistemic and aleatoric uncertainty for my bayesian neural network as described in 3.1 here https://arxiv.org/pdf/2204.09308.pdf

Could you please provide some guidance on how to obtain the mean and variance used by the final layer to generate the output samples? In this way one could estimate epistemic uncertainty for regression tasks.

Thank you!