Adding trainable parameters to each parameter of another model

1

I have a model given by

class M2(nn.Module):
    def __init__( self, dim_x=28*28, dim_y=10):
        super(M2, self).__init__()
        self.phi_y = nn.Sequential(nn.Linear(dim_x, 500), nn.ReLU(inplace=True),
                                   nn.Linear(500, 500), nn.ReLU(inplace=True),
                                   nn.Linear(500, dim_y))

    def forward(self, x):
        pass

I want to add trainable parameters to every parameter in self.phi_y. Then, I would like to use the resulting new network, and use it to make prediction and calculate loss in the forward function. During training, I would like to optimize the parameters of self.phi_y as well as the added parameters.

Thanks!