How could I create a module with learnable parameters

5

Have a look here. *.parameters() creates a generator, normally we would do torch.optim.Adam(loss.parameters(),lr = ...) when dealing with just one set of parameters, but here since you have 2 sets you will need to make a list out of one generator and extend it:

params = list(net.parameters())
params.extend(list(loss.parameters()))
opt = torch.optim.Adam(params,lr=1e-3,weight_decay=5e-4)
3 Likes