add comment

mosaicml · Oct 12, 2023 · 0dae7a7 · 0dae7a7
1 parent ebfacf6
commit 0dae7a7
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
@@ -71,6 +71,8 @@ def __init__(
         if mlflow_logging_config is None:
             mlflow_logging_config = {}
         if self.mlflow_registered_model_name is not None:
+            # Both the metadata and the task are needed in order for mlflow
+            # and databricks optimized model serving to work
             if 'metadata' not in mlflow_logging_config:
                 mlflow_logging_config['metadata'] = {
                     'task': 'llm/v1/completions'