Tensorflow Node exits "Killed"

When training my net, for some reason, often times it fails and I can only see either in the console or dmesg logs that the process was “Killed”.

Is there any mechanism as in python version there is onEarlyEnd (or a similar name) callback that I can pass to my fitDataset function so that the model is saved with the lastest weights?

I fixed this by running tf.tidy(fn) where fn is the function that does most of the computation. So happy it got fixed.