I think your sparse_optimizer contains parameters which were never used in the “amp forward pass”.
If you can verify that that’s the expected case, you should be able to call sparse_optimizer.step() directly without using the GradScaler object.