Hey!
If your data is small enough, it can be completely stored on GPU itself for computation.
data = data.to('cuda')
Else, you can use pin_memory=True to cache the data.
Please refer to this Q&A
Hope this helps.
1 Like