train_generic only accepts model_fn....

would be cool to have a config that accepts an Estimator as input