pytorch
`enumerate()` index doesn't reflect StatefulDataLoader's actual position after loading state
🐛 Describe the bug
Description
When using enumerate() with a StatefulDataLoader that has been restored from a saved state, the enumeration index always starts from 0, even though the dataloader correctly continues from its saved position. This causes issues for any logic that depends on the enumeration index.
To solve this, it might be helpful to add a custom enumerator method on StatefulDataLoader that tracks the actual batch position, or at least document this behavior in the "Stateful DataLoader Tutorial" to help users understand this limitation and find workarounds.
Reproduction
import torch
from torch.utils.data import TensorDataset
from torchdata.stateful_dataloader import StatefulDataLoader
# Create a simple dataset with 10 samples
dataset = TensorDataset(torch.arange(10))
dataloader = StatefulDataLoader(dataset, batch_size=2, shuffle=False)
print("Initial iteration:")
for i, (batch,) in enumerate(dataloader):
print(f"enumerate index: {i}, batch data: {batch}")
if i == 2: # Stop after 3 batches
# Save state after processing batches 0, 1, 2
state = dataloader.state_dict()
break
print("\nCreating new dataloader and loading state:")
dataloader2 = StatefulDataLoader(dataset, batch_size=2, shuffle=False)
dataloader2.load_state_dict(state)
print("\nResuming iteration:")
for i, (batch,) in enumerate(dataloader2):
print(f"enumerate index: {i}, batch data: {batch}")
print("\nPROBLEM: enumerate index starts from 0, but batch data continues from [6, 7]")
print("This breaks any logic that depends on the enumerate index")
Output
Initial iteration:
enumerate index: 0, batch data: tensor([0, 1])
enumerate index: 1, batch data: tensor([2, 3])
enumerate index: 2, batch data: tensor([4, 5])
Creating new dataloader and loading state:
Resuming iteration:
enumerate index: 0, batch data: tensor([6, 7])
enumerate index: 1, batch data: tensor([8, 9])
PROBLEM: enumerate index starts from 0, but batch data continues from [6, 7]
This breaks any logic that depends on the enumerate index
Expected behavior
The enumerate index should be 3, 4 when resuming, matching the actual batch position in the dataset.
Actual behavior
The enumerate index starts from 0, even though the dataloader correctly continues from batch 3.
Impact
This breaks logic that relies on the batch index within an epoch.
Versions
Collecting environment information... PyTorch version: 2.6.0+cu124 Is debug build: False CUDA used to build PyTorch: 12.4 ROCM used to build PyTorch: N/A
OS: Ubuntu 22.04.4 LTS (x86_64) GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0 Clang version: 14.0.0-1ubuntu1.1 CMake version: version 3.31.6 Libc version: glibc-2.35
Python version: 3.11.13 (main, Jun 4 2025, 08:57:29) [GCC 11.4.0] (64-bit runtime) Python platform: Linux-6.1.123+-x86_64-with-glibc2.35 Is CUDA available: False CUDA runtime version: 12.5.82 CUDA_MODULE_LOADING set to: N/A GPU models and configuration: Could not collect Nvidia driver version: Could not collect cuDNN version: Probably one of the following: /usr/lib/x86_64-linux-gnu/libcudnn.so.9.2.1 /usr/lib/x86_64-linux-gnu/libcudnn_adv.so.9.2.1 /usr/lib/x86_64-linux-gnu/libcudnn_cnn.so.9.2.1 /usr/lib/x86_64-linux-gnu/libcudnn_engines_precompiled.so.9.2.1 /usr/lib/x86_64-linux-gnu/libcudnn_engines_runtime_compiled.so.9.2.1 /usr/lib/x86_64-linux-gnu/libcudnn_graph.so.9.2.1 /usr/lib/x86_64-linux-gnu/libcudnn_heuristic.so.9.2.1 /usr/lib/x86_64-linux-gnu/libcudnn_ops.so.9.2.1 Is XPU available: False HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True
CPU: Architecture: x86_64 CPU op-mode(s): 32-bit, 64-bit Address sizes: 46 bits physical, 48 bits virtual Byte Order: Little Endian CPU(s): 2 On-line CPU(s) list: 0,1 Vendor ID: GenuineIntel Model name: Intel(R) Xeon(R) CPU @ 2.20GHz CPU family: 6 Model: 79 Thread(s) per core: 2 Core(s) per socket: 1 Socket(s): 1 Stepping: 0 BogoMIPS: 4399.99 Flags: fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities Hypervisor vendor: KVM Virtualization type: full L1d cache: 32 KiB (1 instance) L1i cache: 32 KiB (1 instance) L2 cache: 256 KiB (1 instance) L3 cache: 55 MiB (1 instance) NUMA node(s): 1 NUMA node0 CPU(s): 0,1 Vulnerability Gather data sampling: Not affected Vulnerability Itlb multihit: Not affected Vulnerability L1tf: Mitigation; PTE Inversion Vulnerability Mds: Vulnerable; SMT Host state unknown Vulnerability Meltdown: Vulnerable Vulnerability Mmio stale data: Vulnerable Vulnerability Reg file data sampling: Not affected Vulnerability Retbleed: Vulnerable Vulnerability Spec rstack overflow: Not affected Vulnerability Spec store bypass: Vulnerable Vulnerability Spectre v1: Vulnerable: __user pointer sanitization and usercopy barriers only; no swapgs barriers Vulnerability Spectre v2: Vulnerable; IBPB: disabled; STIBP: disabled; PBRSB-eIBRS: Not affected; BHI: Vulnerable Vulnerability Srbds: Not affected Vulnerability Tsx async abort: Vulnerable
Versions of relevant libraries: [pip3] intel-cmplr-lib-ur==2025.2.0 [pip3] intel-openmp==2025.2.0 [pip3] mkl==2025.2.0 [pip3] numpy==2.0.2 [pip3] nvidia-cublas-cu12==12.4.5.8 [pip3] nvidia-cuda-cupti-cu12==12.4.127 [pip3] nvidia-cuda-nvrtc-cu12==12.4.127 [pip3] nvidia-cuda-runtime-cu12==12.4.127 [pip3] nvidia-cudnn-cu12==9.1.0.70 [pip3] nvidia-cufft-cu12==11.2.1.3 [pip3] nvidia-curand-cu12==10.3.5.147 [pip3] nvidia-cusolver-cu12==11.6.1.9 [pip3] nvidia-cusparse-cu12==12.3.1.170 [pip3] nvidia-cusparselt-cu12==0.6.2 [pip3] nvidia-nccl-cu12==2.21.5 [pip3] nvidia-nvjitlink-cu12==12.4.127 [pip3] nvidia-nvtx-cu12==12.4.127 [pip3] nvtx==0.2.12 [pip3] optree==0.17.0 [pip3] pynvjitlink-cu12==0.7.0 [pip3] tbb==2022.2.0 [pip3] tcmlib==1.4.0 [pip3] torch==2.6.0+cu124 [pip3] torchao==0.10.0 [pip3] torchaudio==2.6.0+cu124 [pip3] torchdata==0.11.0 [pip3] torchsummary==1.5.1 [pip3] torchtune==0.6.1 [pip3] torchvision==0.21.0+cu124 [pip3] triton==3.2.0 [pip3] umf==0.11.0 [conda] Could not collect
I have noticed this as well. If we are saving model and optimizer states based on the number of steps and stateful dataloader resets back to index 0, won't this be a potential issue?
Hi @divyanshk ,
Do you have any input on how to resolve this issue?
I greatly appreciate your time.
Thank you,
Enrico
@yuvalatzmon I have added a custom enumerate method in #1505 LMK what you think
Hi,
I do greatly appreciate all of your work on this.
Although I was wondering when the custom enumerate method will be merged into the package.
Thank you,
Enrico