Hi @kermitt2
It might be useful to provide an example Jupyter Notebook.
I have created an example (although currently using my "extension" of DeLFT to make it work):
That could be extended to for example define a custom model.