Fix SAETrain dead lock (with segement fault)