Automatic Mixed Precision training (via torch.cuda.amp) can be used to speed up your training and I think the best place to get started is to check out the docs, the Automatic Mixed Precision recipe as well as the examples 1.