Pad chunks to make batchable

pytorch · Feb 6, 2025 · d146017 · d146017
1 parent 0acdf05
commit d146017
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/examples/ibm_rescaling/rescaling_demo.py b/examples/ibm_rescaling/rescaling_demo.py
@@ -73,6 +73,8 @@
 
 # Build dataloader
 data = ScalableReader(datapath, rank, world_size, ArrowHandler, -1, seed=args.seed, max_chunksize=30, n_logical_shards=args.logical_shards)
+# Pad entries to make them batch-able
+data = PreprocessDataset(data, lambda x: x + [-1]*(30-len))
 # Statelessly convert all outputs to tensors
 data = PreprocessDataset(data, torch.tensor)
 # Wrap in StatefulDataLoader