Weighted average of model parameter with trainable scalar

autograd
1

Hi.
I want to find a variable c in order to make weighted average between two models.

For models(modelA, modelB) which has the same structure, I want to find variable c that minimize loss of the averaged model (modelC)
modelC = (1-c)*modelA + c*modelB

modelA = nn.Linear(2, 1)
modelB = nn.Linear(2, 1)
modelC = nn.Linear(2, 1)


c = nn.Parameter(torch.rand(1))
c.requires_grad = True
print('C:', c)

criterion = nn.CrossEntropyLoss()

for p in modelA.parameters():
    p.requires_grad = False
for p in modelB.parameters():
    p.requires_grad = False
for p in modelC.parameters():
    p.requires_grad = False

for name, param in modelC.named_parameters():
    param = c*modelA.get_parameter(name).data() + (1+c)*modelB.get_parameter(name).data()
    
x = torch.randn(1, 2)
y = modelC(x)
y_hat = torch.randn(1, 1)
loss = criterion(y, y_hat)
loss.backward()
print(c.grad)