Custom dataloader running out of memory


I wrote the following " Dataloader " for my validation dataset, it seems to work just fine as can be seen below

Here is the evaluation function that I am using:

The logic is correct, The problem is that when I call the evaluate_val_accuracy() with my model, I can see that my memory consumption on gpu keeps on increasing and eventually it runs out of memory throwing error like

I am guessing the problem can not be with the model as while training it starts to train, and only throws error after an epoch where I call this evaluation function. The whole point of using yield was so that I don’t have to do .cuda() on all the data at once, but it doesn’t work. any clues?