push

torchtitan/datasets/hf_datasets.py

@@ -9,16 +9,16 @@
 from typing import Any, Callable, Dict, List, Optional
 
 import torch
-
-from datasets import Dataset, load_dataset
-from datasets.distributed import split_dataset_by_node
 from torch.distributed.checkpoint.stateful import Stateful
 from torch.utils.data import IterableDataset
 from torchdata.stateful_dataloader import StatefulDataLoader
 
 from torchtitan.datasets.tokenizer import Tokenizer
 from torchtitan.logging import logger
 
+from datasets import Dataset, load_dataset
+from datasets.distributed import split_dataset_by_node
+
 
 def _load_c4_dataset(dataset_path: str):
     """Load C4 dataset with default configuration."""