ensure the sampler do not goes past the file in the last rank.

argonne-lcf · Aug 31, 2024 · 2321952 · 2321952
1 parent 682a300
commit 2321952
Showing 1 changed file with 4 additions and 2 deletions.
diff --git a/dlio_benchmark/data_loader/torch_data_loader.py b/dlio_benchmark/data_loader/torch_data_loader.py
@@ -92,8 +92,10 @@ def __init__(self, rank, size, num_samples, epochs):
         self.epochs = epochs
         samples_per_proc = int(math.ceil(num_samples/size)) 
         start_sample = self.rank * samples_per_proc
-        end_sample = (self.rank + 1) * samples_per_proc
-        self.indices = list(range(start_sample, end_sample))
+        end_sample = (self.rank + 1) * samples_per_proc - 1
+        if end_sample > num_samples - 1:
+            end_sample = num_samples - 1
+        self.indices = list(range(start_sample, end_sample + 1))
 
 
     def __len__(self):