@Jason , thanks for the reply.
You’re right, reducing the image size fixed the problem.
I was wondering if there is a way to notify the application when we are running out of memory? As I mentioned, checking the tf.memory() at batchEnd doesn’t seem to catch this state since the app crashes in between these callbacks. The warning we currently have is great for development purposes but it would be better to be able to detect this state and let the application decide whether the training should be aborted, rather than allowing the GPU to crash.
Thanks,
A.