[Python] IterableDataset items consistently fail filter in collate_fn on first batch, despite...

I'm encountering a puzzling issue while training a transformer model on soccer event sequences using PyTorch's IterableDataset and a custom collate_fn (potentially within the Hugging Face Trainer, but the core issue seems related to the DataLoader interaction).

My IterableDataset yields dictionaries containing tensors (input_cat, input_cont, etc.). I've added print statements right before the yield statement, confirming that valid dictionaries with the expected tensor keys and shapes are being produced.

The DataLoader collects these items (e.g., batch_size=16). However, when the list of collected items reaches my collate_fn, a filter check at the beginning removes all items from the batch. This happens consistently on the very first batch of training.

The filter check is: batch = [b for b in batch if isinstance(b, dict) and "input_cat" in b]

Because this filter removes all items, the collate_fn then detects len(batch) == 0 and returns a signal to skip the batch ({"skip_batch": True}). The batch received by collate_fn is a list of 16 empty dictionaries.

Additionally, batch size is 16 and block size is 16.

The code is as follows:

class IterableSoccerDataset(IterableDataset):
def __init__(self, sequences: List[List[Dict]], idx: FeatureIndexer, block_size: int, min_len: int = 2):
super().__init__()
self.sequences = sequences
self.idx = idx
self.block_size = block_size
self.min_len = min_len
self.pos_end_cat = np.array([idx.id_for("event_type", idx.POS_END) if col=="event_type" else 0
for col in ALL_CAT], dtype=np.int64)
self.pos_end_cont = np.zeros(len(ALL_CONT), dtype=np.float32)
print(f"IterableSoccerDataset initialized with {len(sequences)} sequences.")

def __iter__(self) -> Iterator[Dict[str, torch.Tensor]]:
rng = np.random.default_rng()
for seq in self.sequences:
if len(seq) < self.min_len:
continue

# encode
cat, cont = [], []
for ev in seq:
c, f = self.idx.encode(pd.Series(ev))
cat.append(c)
cont.append(f)
cat.append(self.pos_end_cat)
cont.append(self.pos_end_cont)

cat = np.stack(cat) # (L+1,C)
cont = np.stack(cont) # (L+1,F)
L = len(cat) # includes POS_END

# decide window boundaries
if L <= self.block_size + 1:
starts = [0] # take the whole thing
else:
# adaptive stride: roughly 50 % overlap
stride = max(1, (L - self.block_size) // 2)
starts = list(range(0, L - self.block_size, stride))
# ensure coverage of final token
if (L - self.block_size) not in starts:
starts.append(L - self.block_size)

print(L, len(starts))

for s in starts:
e = min(s + self.block_size + 1, L)
inp_cat = torch.from_numpy(cat[s:e-1]) # length ≤ block
tgt_cat = torch.from_numpy(cat[s+1:e])
inp_cont = torch.from_numpy(cont[s:e-1])
tgt_cont = torch.from_numpy(cont[s+1:e])

print(f"DEBUG: Yielding item - input_cat shape: {inp_cat.shape}, seq_len: {inp_cat.size(0)}")

yield {
"input_cat": inp_cat,
"input_cont": inp_cont,
"tgt_cat": tgt_cat,
"tgt_cont": tgt_cont,
}


def collate_fn(batch):
batch = [b for b in batch
if isinstance(b, dict) and "input_cat" in b]

if len(batch) == 0:
return {"skip_batch": True}

# ... rest of code


I have tried:

1. 
   Successful yielding - confirmed via prints that the iter method does yield dictionaries with the key "input_cat" and others, containing tensors.
   
   
2. 
   collate_fn receives items - confirmed via prints that collate_fn receives a list (batch) with the correct number of items (equal to batch_size).
   
   
3. 
   Filtering failure - the specific filter isinstance(b, dict) and "input_cat" in b evaluates to False for every item received by collate_fn in that first batch (as they are all just empty dictionaries).
   
   
4. 
   num_workers - I suspected this might be related to multiprocessing (dataloader_num_workers > 0), potentially due to serialization/deserialization issues between workers and the main process. However, did not make a difference when I set dataloader_num_workers=0.
   

What could cause items that appear correctly structured just before being yielded by the IterableDataset to consistently fail the isinstance(b, dict) and "input_cat" in b check when they arrive as a list in the collate_fn, especially on the very first batch? I am at a loss for what to do.

Many thanks!

Continue reading...