I am trying to train a model that uses pixels as data samples.
Because the dataset is small, i am trying to keep the entire dataset on the GPU.
When timing the run of the model, the first 5 iterations are fast, but then it slows down drastically:

The first iterations is slow due to “cold start”, then the next 4 are fast, and the next ones are vary slow.
If i keep the data on the CPU, and move the batch to the GPU one batch at a time, the model runs at full speed all the time, but the time it takes to move the tensors is too long.
If i keep the data on the GPU, but after every iteration i call torch.cuda.empty_cache() the model runs at full speed, but empty_cache takes 1.25 seconds …
I tried to detach and clone the batch before using the forward pass, it did nothing.
what could be the cause of this problem. and how could it be fixed ?
I would appreciate any advice