I don’t know which “module” you mean and am still unsure how you’ve profiled it and narrowed down the bottleneck.
In any case, to preload the data, just push it to the device:
data = data.to(device)

and wrap it into e.g. a TensorDataset.



 AP_M:

Can the pytorch NN with batch size of 1 and big dataset be used efficiently with GPU’s?


It depends on the model. If the GPU workload is tiny, your script might suffer from the kernel launches and general CPU overhead.