В 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

Но я просто могу вычислить град и не могу их обновить. Я думаю, причина в том, что я не добавляю их в параметр оптимизатора.

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

loss.backward()
optimizer.step()

Итак, я хочу спросить, как сделать гамма- и бета-обновление на этапе поезда? Спасибо тебе!