Hi @chenglu, thanks for the answer. I think that training and testing batch size matters. I tried out with larger test batch size and it didn’t work, but when I reduce the test batch size, it works. So it seems that even if the dataloader is emptied out at the end of every epoch, it occupies memory.