From 8f595fea755dc15fe8a10bdd6eafb9f7f7cb9dfc Mon Sep 17 00:00:00 2001 From: dodatboii Date: Wed, 14 Oct 2026 16:10:57 +0800 Subject: [PATCH] fix: align micro-batch metadata with payload slices Signed-off-by: dodatboii --- transfer_queue/dataloader/streaming_dataset.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/transfer_queue/dataloader/streaming_dataset.py b/transfer_queue/dataloader/streaming_dataset.py index e308ef6a..7f5675a8 100644 --- a/transfer_queue/dataloader/streaming_dataset.py +++ b/transfer_queue/dataloader/streaming_dataset.py @@ -331,7 +331,6 @@ def chunk_batch_fn(td, batch_meta, micro_batch_size=1): # Calculate number of splits (handles uneven division) num_splits = (total_size + micro_batch_size - 1) // micro_batch_size splits = [] - batch_meta_list = batch_meta.chunk(num_splits) # Pre-process: separate nested tensors so that native ``td[start:end]`` works. # TensorDict delegates slice to each field; nested tensors do not support @@ -360,6 +359,6 @@ def chunk_batch_fn(td, batch_meta, micro_batch_size=1): for key, (tensors, layout) in nested_fields.items(): chunk_td[key] = torch.nested.as_nested_tensor(tensors[start:end], layout=layout) - splits.append((chunk_td, batch_meta_list[i])) + splits.append((chunk_td, batch_meta[start:end])) return splits