If you could trace the model, all operations should be recorded (thus also the transformations) and amp should work. This would also mean that you don’t need to run the traced model in an autocast region anymore.
However, scripted models are currently still WIP.