Add custom collate function for Getting Started example (resolves the collate_fn TypeError) (#607)

* fix: getting started default example by adding custom collate function for StreamingDataLoader

* [pre-commit.ci] auto fixes from pre-commit.com hooks

for more information, see https://pre-commit.ci

---------

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>

Author: bhimrazy

README.md

@@ -127,11 +127,19 @@ Load the data by replacing the PyTorch DataSet and DataLoader with the Streaming
 ```python
 import litdata as ld
 
-train_dataset = ld.StreamingDataset('s3://my-bucket/fast_data', shuffle=True, drop_last=True)
-train_dataloader = ld.StreamingDataLoader(train_dataset)
+dataset = ld.StreamingDataset('s3://my-bucket/fast_data', shuffle=True, drop_last=True)
 
-for sample in train_dataloader:
-    img, cls = sample['image'], sample['class']
+# Custom collate function to handle the batch (Optional)
+def collate_fn(batch):
+    return {
+        "image": [sample["image"] for sample in batch],
+        "class": [sample["class"] for sample in batch],
+    }
+
+
+dataloader = ld.StreamingDataLoader(dataset, collate_fn=collate_fn)
+for sample in dataloader:
+    img, cls = sample["image"], sample["class"]
 ```
 
 **Key benefits:**

examples/getting_started/README.md

@@ -64,5 +64,16 @@ img = sample['image']
 cls = sample['class']
 
 # Create dataLoader and iterate over it to train your AI models.
-dataloader = StreamingDataLoader(dataset)
+
+# Custom collate function to handle the batch (Optional)
+def collate_fn(batch):
+    return {
+        "image": [sample["image"] for sample in batch],
+        "class": [sample["class"] for sample in batch],
+    }
+
+
+dataloader = StreamingDataLoader(dataset, collate_fn=collate_fn)
+for sample in dataloader:
+    img, cls = sample["image"], sample["class"]
 ```

examples/getting_started/stream.py

@@ -12,4 +12,16 @@
 cls = sample["class"]
 
 # Create dataLoader and iterate over it to train your AI models.
-dataloader = StreamingDataLoader(dataset)
+
+
+# Custom collate function to handle the batch (Optional)
+def collate_fn(batch):
+    return {
+        "image": [sample["image"] for sample in batch],
+        "class": [sample["class"] for sample in batch],
+    }
+
+
+dataloader = StreamingDataLoader(dataset, collate_fn=collate_fn)
+for sample in dataloader:
+    img, cls = sample["image"], sample["class"]