Add weight keyword arg to train fn for unbalanced training data