bugfix: train - fix learn_mask training