From 0dae7a762a7e3966cafc6b011d6a1500bdc1e049 Mon Sep 17 00:00:00 2001
From: Daniel King <daniel@mosaicml.com>
Date: Wed, 11 Oct 2023 17:47:01 -0700
Subject: [PATCH] add comment

---
 llmfoundry/callbacks/hf_checkpointer.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/llmfoundry/callbacks/hf_checkpointer.py b/llmfoundry/callbacks/hf_checkpointer.py
index 843ca102f3..a02f8b98a0 100644
--- 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'