TypeError: log_model() got an unexpected keyword argument 'task'

Question

This is really stumping me. I have MLflow version 2.5.0, and here is the relevant code:

torch_version = torch.__version__.split("+")[0]
with mlflow.start_run() as run:  
    mlflow.pyfunc.log_model(
        python_model=Pythia(),
        artifacts={'repository' : $PATH},
        pip_requirements=[f"torch=={torch_version}", 
                          f"transformers=={transformers.__version__}", 
                          f"accelerate=={accelerate.__version__}", "einops", "sentencepiece"],
        input_example=input_example,
        task="text2text-generation",
        signature=signature
    )

And here is the stacktrace:

TypeError                                 Traceback (most recent call last)
File <command-853386702488317>, line 144
    142 torch_version = torch.__version__.split("+")[0]
    143 with mlflow.start_run() as run:  
--> 144     mlflow.pyfunc.log_model(
    145         python_model=Pythia(),
    146         artifacts={'repository' : $PATH},
    147         pip_requirements=[f"torch=={torch_version}", 
    148                           f"transformers=={transformers.__version__}", 
    149                           f"accelerate=={accelerate.__version__}", "einops", "sentencepiece"],
    150         input_example=input_example,
    151         task="text2text-generation",
    152         signature=signature
    153     )
    155 import mlflow
    156 mlflow.set_registry_uri("databricks-uc")

TypeError: log_model() got an unexpected keyword argument 'task'

I know I don't have a random variable also called task, so I'm really just at a loss. I'm looking at the documentation and it shows task being used like how I'm using it above, so color me confused. I am using Databricks, g5.4xl, DBR 14.0 ML.