Improve out of memory handling