Description & Motivation

Right now we can track models on mlflow for all tasks that are lightning based, but we can't do it for sklearn based classification and patch classification.

It would be a nice addition to have this feature also for these two tasks as it will increase the amount of traceability of our models