I’m unsure if there is an easy way to “tune” the autocast list (@mcarilli might know), but for now you could disable autocast for the desired operations, transform your data to bfloat16() and call the operations directly.