Merge branch 'main' into feat/group-size-divisibility-check

HDCharles · web-flow · commit a85395b24fc1 · 2026-02-16T10:32:21.000-05:00
diff --git a/src/llmcompressor/datasets/utils.py b/src/llmcompressor/datasets/utils.py
@@ -260,7 +260,7 @@ def _make_sampler(args: DatasetArguments, dataset: Dataset) -> Sampler:
         start, end = _get_partition_start_end(
             num_samples, dist.get_rank(), dist.get_world_size()
         )
-        dataset = dataset[start:end]
+        dataset = dataset.select(range(start, end))
 
     if num_samples is not None and num_samples > len(dataset):
         logger.warning(

Original file line number	Diff line number	Diff line change
`@@ -260,7 +260,7 @@ def _make_sampler(args: DatasetArguments, dataset: Dataset) -> Sampler:`
`260`	`260`	`start, end = _get_partition_start_end(`
`261`	`261`	`num_samples, dist.get_rank(), dist.get_world_size()`
`262`	`262`	`)`
`263`		`- dataset = dataset[start:end]`
	`263`	`+ dataset = dataset.select(range(start, end))`
`264`	`264`
`265`	`265`	`if num_samples is not None and num_samples > len(dataset):`
`266`	`266`	`logger.warning(`