One thing I've noticed is if something goes during initialization of the training (wether it'd be OOM or something else), I always get the same "...task failed" error. No mention that it ran out of memory even if it's clearly the case (using lower batch size resolves the issue)