AMP for two optimizers

mixed-precision
1

Hi, all.
I want to try apply AMP for the two optimizers .
I just added all losses to total_loss.
Below codes are correctly optimize the two optimizers?

optimizer.zero_grad()
optimizer2.zero_grad()
scaler.scale(total_loss).backward()
scaler.step(optimizer)
scaler.step(optimizer2)
scaler.update()

Thanks in advance.

2

Yes, your code should work. The docs also give an example which fits your use case.