I would stick to the recommended usage of moving the model’s forward and loss calculation into the autocast context. These pieces from the docs might be interesting to check your use case:

autocast should wrap only the forward pass(es) of your network, including the loss computation(s). Backward passes under autocast are not recommended. Backward ops run in the same type that autocast used for corresponding forward ops.


Floating-point Tensors produced in an autocast-enabled region may be float16. After returning to an autocast-disabled region, using them with floating-point Tensors of different dtypes may cause type mismatch errors. If so, cast the Tensor(s) produced in the autocast region back to float32 (or other dtype if desired). If a Tensor from the autocast region is already float32, the cast is a no-op, and incurs no additional overhead.

Besides that autocast uses an internal cache to avoid re-casting constant tensors, which would never be freed if you are wrapping the entire code into an autocast region and you might thus have to either disable it or could suffer from increasing memory usage.