I want to make a loss module with trainable parameter, so I made the following:
class CostumLoss(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.gamma = torch.nn.Parameter(torch.FloatTensor([.5]))

    def forward(self, od_loss, depth_loss):
        print(self.gamma)
        loss = od_loss + self.gamma * depth_loss
        print(self.gamma)
        return loss

c_loss = CostumLoss()
c_loss.train()
total_loss = c_loss(od_loss , depth_loss)
total_loss.backward()

but printing the learnable parameter gamma tells that the parameter didn’t change.