Bayesian Neural Network for Nonlinear Regression Analysis

Hello friends, I want to implement nonlinear regression analysis given to Bayesian neural network. But I’m a beginner and I found that the output curve is a straight line while using the linear regression code in the tutorial. I would like to ask how to define a nonlinear network, is there any related function? Or do I need to write it myself? Looking forward to your reply, thanks!

class BayesianRegression(PyroModule):
def init(self, in_features, out_features):
super().init()
self.linear = PyroModule[nn.Linear](in_features, out_features)
self.linear.weight = PyroSample(dist.Normal(0., 1.).expand([out_features, in_features]).to_event(2))
self.linear.bias = PyroSample(dist.Normal(0., 10.).expand([out_features]).to_event(1))

def forward(self, x, y=None):
    sigma = pyro.sample("sigma", dist.Uniform(0., 10.))
    mean = self.linear(x).squeeze(-1)
    with pyro.plate("data", x.shape[0]):
        obs = pyro.sample("obs", dist.Normal(mean, sigma), obs=y)
    return mean