`enumerate()` index doesn't reflect StatefulDataLoader's actual position after loading state

[yuvalatzmon]

🐛 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

[conceptofmind]

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?

[conceptofmind]

Hi @divyanshk ,

Do you have any input on how to resolve this issue?

I greatly appreciate your time.

Thank you,

Enrico

[ramanishsingh]

@yuvalatzmon I have added a custom enumerate method in #1505
LMK what you think

[conceptofmind]

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