Your autocast example looks correct.



 tree33:

And from my knowledge, gradients are scaled during mixed precision


Yes, if you are using a GradScaler object.



 tree33:

If I have to write my own backward function for sparse conv layers (warped in autocast(disabled)),
Do I still need to consider the scale?


No, gradient scaling and auto casting are working together but are independent from each other.
You don’t need to explicitly implement gradient scaling for the custom layer.