Solved by carloalbertobarbano in post #2 
              
              
                Is the question about how to make A and B learnable parameters? 
In that case, you could add A and B to your model class definition as nn.Parameter and then compute the loss as 
loss = model.A * bce(output, gt) + model.B