在pytorch中，如何为损失的不同重量部分定义优化器或损失函数？

Question

我定义了一个包含不同部分的混合损失函数。我想训练不同部位的重量参数。

class Loss(nn.Module):
    def __init__(self):
        super(Loss, self).__init__()
        self.gamma = torch.nn.Parameter(torch.FloatTensor(1), requires_grad=True).cuda()
        self.gamma.data.fill_(0.1)
        self.beta = torch.nn.Parameter(torch.FloatTensor(1), requires_grad=True).cuda()
        self.beta.data.fill_(0.1)

    def forward(self, data1, label1, data2, label2):
        loss1 = torch.mean(data1 - label1)
        loss2 = torch.mean(abs(data2 - label2))

        loss = loss1 * self.gamma + loss2 * self.beta
        return loss

但是我只能计算出grad，无法更新它们。我认为原因是我没有将它们添加到优化程序参数中。

optimizer = optim.Adam(model.parameters(), lr=args.lr_initial)

loss.backward()
optimizer.step()

所以我想问一下在训练阶段如何进行伽玛和Beta更新？谢谢你！