fixed docstring

mosaicml · Feb 7, 2024 · 4dd8524 · 4dd8524
1 parent 165c818
commit 4dd8524
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/llmfoundry/callbacks/curriculum_learning_callback.py b/llmfoundry/callbacks/curriculum_learning_callback.py
@@ -26,7 +26,7 @@ class CurriculumLearning(Callback):
         dataset_index (int): The index of the dataset currently being used.
         current_dataset_config (Dict): The configuration of the dataset currently
             being used.
-        dataset (StreamingDataset): The dataset currently being used.
+        dataloader (Union[DataLoader, Iterable]): The training dataloader currently being used.
     """
 
     def __init__(self, dataset_index: int, dataloader: Union[DataLoader,