@moreshud as you can see in the last line if I reduce the batch size to 256, then my GPU can handle it, but my problem here is it reaches out of memory in the second epoch. If there is not enough Memory for training it should raise the bug from the first epoch, am I right?