Thanks for the reply, @lgusm .
I don’t think this is a memory issue. I purposefully reduced the number of images down to 64 so that such a problem can’t happen. I also checked my task manager and didn’t see any problem there either.
Could there be a problem with the tensorflow version that I’m using? It’s not likely the problem, but I think I have an earlier version than the one used on colab.