如何使用pytorch-lightning将模型权重保存到mlflow跟踪服务器？

pip install pytorch-lightning mlflow

import torch
from pytorch_lightning import Trainer
from pytorch_lightning.loggers import MLFlowLogger

class MyModel(pl.LightningModule):
    def __init__(self):
        super().__init__()
        # 模型定义

    def training_step(self, batch, batch_idx):
        # 训练逻辑

    def configure_optimizers(self):
        # 优化器定义

model = MyModel()
trainer = Trainer(logger=MLFlowLogger())
trainer.fit(model)

mlflow.set_tracking_uri('http://your_mlflow_server')
mlflow.set_experiment('your_experiment_name')

class SaveModelCallback(pl.Callback):
    def on_epoch_end(self, trainer, pl_module):
        mlflow.pytorch.log_model(pl_module, "model")

trainer = Trainer(callbacks=[SaveModelCallback()], logger=MLFlowLogger())
trainer.fit(model)