Model_T and Model_Y in DML/ORF

Hi! I have a question on model_T and model_Y in DML/ORF. I notice that, by default, model_T and model_Y use Lasso in scikit-learn, while I'm thinking to handle the high-dimensional confounding factors, W, in a nonlinear way, such as using a neural network to fit model_T and model_Y. Theoretically, it's workable (just as mentioned in papers). But in practice, is that possible to train a neural network in Keras and then use it to fit model_T and model_Y?

I try to implement this idea, but get some problems. After trained my neural network in Keras, I saved it using joblib.dump, and then use it to fit model_T and model_Y while calling ORF. However, Python returns some errors:
AttributeError: 'History' object has no attribute 'predict'
and
joblib.externals.loky.process_executor.BrokenProcessPool: A result has failed to un-serialize. Please ensure that the objects returned by the function are always picklable.

I wonder whether Keras models can be used to fit model_T and model_Y? It looks like simply dumping the model from Keras doesn't work. Maybe I need to write a class, like those in scikit-learn, to make Keras models compatible with ORF/DML? Any ideas? Thanks!

Start at the ORF/DML entry points for model_T and model_Y and compare their estimator expectations with the Keras model handling described in the issue. Reproduce the joblib-loaded model errors; done should establish whether Keras models are supported and, if not, document the required compatibility path or scope a change.