Create an EasyGuide for a PyroModule model, SVI

Hello, I am new to pyro and I have an issue to define an EasyGuide for a Pyro module. I have a deep learning model define with a pyro module, inspired from the tutorial Tutorial 1: Bayesian Neural Networks with Pyro — UvA DL Notebooks v1.2 documentation

class mymodel(PyroModule):
def init(self, in_dim=1, out_dim=1, hid_dim=4, n_hid_layers=2):
super().init()
self.activation = nn.GELU() # could also be ReLU

    # layer sizes : 
    self.layer_sizes = [in_dim] + n_hid_layers * [hid_dim] + [out_dim]
    layer_list = [PyroModule[nn.Linear](self.layer_sizes[idx], self.layer_sizes[idx-1]) for idx in
                  range(1, len(self.layer_sizes))]
    #define the neural network
    self.layers = PyroModule[torch.nn.ModuleList](layer_list)

    for layer_idx, layer in enumerate(self.layers):
        layer.bias = PyroSample(dist.Normal(0,1).expand([self.layer_sizes[layer_idx + 1]]).to_event(1))
        layer.weight = PyroSample(dist.Normal(0,1).expand(
            [self.layer_sizes[layer_idx + 1], self.layer_sizes[layer_idx]]).to_event(2))


def forward(self, x, y=None):
    x=x.reshape(1,-1)
    x = self.activation(self.layers[0](x))
    for layer in (self.layers[1:-1]):
        x = self.activation(layer(x))  # hidden --> hidden
    mu = self.layers[-1](x).squeeze()  # hidden --> output
    
    sigma = pyro.sample("sigma", dist.Gamma(.5, 1))  # infer the response noise
    #add observation noise around the predicted mean

#given that i use this neural network for online learning my pyro.plate only needs to be of size 1
    with pyro.plate("data", 1):
        pyro.sample("obs", dist.Normal(mu, sigma * sigma), obs=y)
    return mu

But instead of defining my guide with a simple :
model = mymodel()
guide = AutoDiagonalNormal(model)
I want to add some liberty to my guide defining an EasyGuide, like this for instance :

class GUIDE(EasyGuide):
def init(self, model, in_dim=1, out_dim=1, hid_dim=4, n_hid_layers=5):
super().init(model)
self.activation = nn.GELU()
self.layer_sizes = [in_dim] + n_hid_layers * [hid_dim] + [out_dim]
layer_list = [PyroModule[nn.Linear](self.layer_sizes[idx], self.layer_sizes[idx-1]) for idx in range(1, len(self.layer_sizes))]
self.layers = PyroModuletorch.nn.ModuleList

    for layer_idx, layer in enumerate(self.layers):
        weight_scale = pyro.param('weight_scale'+ str(layer_idx), torch.tensor([1. for _ in range(len(self.layers))]), constraint=constraints.positive)
        weight_loc = pyro.param('weight_loc' + str(layer_idx), torch.tensor([0. for _ in range(len(self.layers))]))
        bias_scale = pyro.param('bias_scale'+ str(layer_idx), torch.tensor([1. for _ in range(len(self.layers))]), constraint=constraints.positive)
        bias_loc = pyro.param('bias_loc'+ str(layer_idx), torch.tensor([0. for _ in range(len(self.layers))]))

        layer.weight = PyroSample(dist.Normal(weight_loc[layer_idx], weight_scale[layer_idx]).expand([self.layer_sizes[layer_idx+1], self.layer_sizes[layer_idx]]).to_event(2))
        layer.bias = PyroSample(dist.Normal(bias_loc[layer_idx], bias_scale[layer_idx]).expand([self.layer_sizes[layer_idx + 1]]).to_event(1))

def guide(self, x, y=None):
    x = x.reshape(1, -1)
    x = self.activation(self.layers[0](x))  # input --> hidden
    for temp, layer in enumerate(self.layers[1:-1]):
        x = self.activation(layer(x))
    mu = self.layers[-1](x).squeeze() 
    sigma = pyro.sample("sigma", dist.Gamma(0.5, 1))  # infer the response noise

    with pyro.plate("data", x.shape[0]):
        pyro.sample("obs", dist.Normal(mu, sigma * sigma), obs=y)
    return mu
        
def forward(self, x, y=None):
    x = x.reshape(1, -1)
    x = self.activation(self.layers[0](x))  # input --> hidden
    for temp, layer in enumerate(self.layers[1:-1]):
        x = self.activation(layer(x))
    mu = self.layers[-1](x).squeeze() 
    sigma = pyro.sample("sigma", dist.Gamma(0.5, 1))  # infer the response noise

    with pyro.plate("data", x.shape[0]):
        pyro.sample("obs", dist.Normal(mu, sigma * sigma), obs=y)
    return mu

But to do so i would need to name the parameters in order to have the same PyroSample in the model and in the guide…
But if I try to name my parameters in my model with the line “layer.weight.name = “name” I have the error message : AttributeError: attribute ‘name’ of ‘torch._C._TensorBase’ objects is not writable”
and if I add the argument name = “name” inside the Pyro Sample I’m being told that the argument name shouldn’t be initialized here :

TypeError: PyroSample.new() got an unexpected keyword argument ‘name’

How should I do? Should I find another way to define my EasyGuide or is there another way to define the name od the Parameters of a PyroSample ? It seems that in Pyromodule I must use PyroSample, that raises a pyro.sample statement but I don’t understand why is that? What is the particularity of a PyroSample and why can’t i directly do a pyro.sample to define the weights of my neural network ? Thank you very much.

NB : If I define many times the same params in the guide with a pyro.param it is because in a next step I won’t initialize them with the same value, i will do Lucun initialisation for instance.