How to train a tensor of tensors in PyTorch

Question

I have a class model:

class Model(nn.Module)

That has 2 learnable parameters:

self.a = torch.nn.Parameter(torch.rand(1, requires_grad=True))
self.b = torch.nn.Parameter(torch.rand(1, requires_grad=True))

It also has a neural network class inside:

class Net_x(nn.Module)

So in order to train all the parameters I combine the learnable parameters "a, b" and the NN's params (under the "init" of class Model):

self.net_x = self.Net_X()
self.params = list(self.net_x.parameters()) 
self.params.extend(list([self.a, self.b]))

This works well and the network is training just fine.

My issue is that now I'm trying to change one of the parameters to be a tensor of tensors of parameters:

self.a = torch.tensor( [[ torch.nn.Parameter(torch.rand(1, requires_grad=True)) ] for i in range(5)] )

That is because at every time step (not epoch), I need to use a different parameter from self.a

Example:

for epoch in range(n_epochs): 
    if timestep<5: 
        val = 50 - a[0] * b 
        loss = 10 - val 
    elif timestep >=5 and timestep < 10: 
        val = 50 - a[1] * b 
        loss = 10 - val

The model runs without issues, but the parameters are not being updated (ie they stay the same at every epoch).

PS I would have added my code, but it's really long. I'm hoping the answer is simple (and if not I'll try to reduce my code and attach it)

Answer 1

Try this:

self.a = torch.tensor( [[ torch.nn.Parameter(torch.rand(1, requires_grad=True)) ] for i in range(5)], requires_grad = True )

requires_grad=False by default so I think that might be the issue