torch.cuda.amp.autocast will use mixed-precision training and cast necessary tensors under the hood for you.
From the docs:
When entering an autocast-enabled region, Tensors may be any type. You should not call
.half()on your model(s) or inputs when using autocasting.