Hi,
I don’t know how complex your data are but take into account that most of the pytorch gens (torch.rand for example) can directly allocate the data into the gpu.
That way you save the time required to move data from cpu to gpu.
You may find this useful How to measure time in PyTorch
This way you can profile cpu vs gpu