I would like to save model weights to mlflow tracking using pytorch-lightning. pytorch-lightning supports logging. However, it seems that saving model weights as a artifact on mlflow is not supported.

At first, I planed to override ModelCheckpoint class to do it, but I found it is difficult for me because of complex Mixin operations.

Anybody knows simple way to accomplish it?