refine

meta-llama · Jan 10, 2025 · 1e915d8 · 1e915d8
1 parent cc21a99
commit 1e915d8
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/llama_stack/providers/inline/post_training/torchtune/datasets/sft.py b/llama_stack/providers/inline/post_training/torchtune/datasets/sft.py
@@ -47,11 +47,10 @@ def __getitem__(self, index: int) -> Dict[str, Any]:
     def _prepare_sample(self, sample: Mapping[str, Any]) -> Dict[str, Any]:
         if self._dataset_type == "instruct":
             sample = llama_stack_instruct_to_torchtune_instruct(sample)
-        elif self._dataset_type == "chat":
+        elif self._dataset_type == "dialog":
             sample = llama_stack_chat_to_torchtune_chat(sample)
         else:
             raise ValueError(f"Invalid dataset type: {self._dataset_type}")
-
         transformed_sample = self._message_transform(sample)
         if "messages" in transformed_sample:
             validate_messages(transformed_sample["messages"])