I would suggest checking out: https://skorch.readthedocs.io

You can write your own linear model with a custom loss function in PyTorch. Then, with skorch use the model in your pipeline.

Limitation: can only work with linear/non-linear classifiers. Can't work with trees.