MLFlow 无法记录 pytorch 模型

Question

After training a model, I am trying to log it to mlflow with mlflow.pytorch.log_model(model, artifact_path="model",pickle_module=pickle) but i get the error:

yaml.representer.RepresenterError: ('cannot represent an object', '1.11.0+cu102')

I definitely send the model to cpu before doing so and confirm its there with next(model.parameters()).device . How can I get mlflow to log my model?

Answer 1

I haven't used MLflow with PyTorch but in the worst case, you can avoid using the PyTorch flavor and just do something like

    with tempfile.TemporaryDirectory() as path:
        model.save(path)
        mlflow.log_artifacts(path, artifact_path="model")