Thank you for your response!
Could you please suggest a way to make the noise dependent on the input, or if you know of any examples here in the discourse or elsewhere on the web I would highly appreciate a link.
Here is the reasoning that I followed when writing this code. The weights and biases of the neural network are random variables. Therefore layer_o (output of the BNN) will already have some variance that depends on the input. My reasoning is that this variance due to the uncertainty in the weights and biases should account for the input dependent noise. But my training data is generated by adding a small constant Gaussian noise at every point (notice the waviness in the training data), so even if the model figured out the exact distributions of weights and biases required to generate this data, it would still need to add a small variance on top to account for the gaussian noise, that is why I am using the sigma = pm.HalfNormal('sigma',sd=0.1). I would appreciate it if you could help me understand the error in my logic.
Edit: By the way, as I expect, the variance in the NN does indeed change with the input (as seen in the top right plot), but the change is very small and is very different from the actual variance in the training data. So another follow up question would be, what is causing that change in the variance if my model is configured to have a fixed noise regardless of the input?