Thank you for the great answer, whole pipeline is a lot clearer now.



 ptrblck:

You should not call half() on the data or the model and apex/amp or torch.cuda.amp.autocast take care of the casting for you.
I’m not sure where the Trainer class comes from, but I assume it’s from a high-level API?
If so, I would expect the same behavior, i.e. no half() calls.


My bad, Trainer comes from PyTorch Lightning, but still sounds like casting will be taken care of because it uses apex/amp under the hood.
I have one followup question:
All my data when in FP16 can fit entirely on the GPU, which to my understanding would be more performant compared to transferring data from host to device even with pin_memory=True and multiple workers.
The problem is that I can never store all the data on the GPU in FP32, even intermediately, because it doesn’t fit. I need to have the half() call applied before the to(cuda_device) / cuda() call. Is there a way for me to move the data to GPU after the half() call that is done for me by amp/apex or torch.cuda.amp? In other words, is there a way for the half() to be done on the CPU tensor for the entire data batch, and have the tensor copied to GPU one-time after?