Missing elements in `map` a batched dataset
### Describe the bug

As outlined [here](https://discuss.huggingface.co/t/length-error-using-map-with-datasets/40969/3?u=sachin), the following collate function drops 5 out of possible 6 elements in the batch (it is 6 because out of the eight, two are bad links in laion). A reproducible [kaggle kernel ](https://www.kaggle.com/sachin/laion-hf-dataset/edit) can be found here.

The weirdest part is when inspecting the sizes of the tensors as shown below, both `tokenized_captions["input_ids"]` and `image_features` show the correct shapes. Simply the output only has one element (with the batch dimension squeezed out).

```python
class CollateFn:        
    def get_image(self, url):
        try:
            response = requests.get(url)
            return Image.open(io.BytesIO(response.content)).convert("RGB")
        except PIL.UnidentifiedImageError:
            logger.info(f"Reading error: Could not transform f{url}")
            return None
        except requests.exceptions.ConnectionError:
            logger.info(f"Connection error: Could not transform f{url}")
            return None
        
        
    def __call__(self, batch):
        images = [self.get_image(url) for url in batch["url"]]
        captions = [caption for caption, image in zip(batch["caption"], images) if image is not None]
        images = [image for image in images if image is not None]
        
        tokenized_captions = tokenizer(
            captions,
            padding="max_length",
            truncation=True,
            max_length=tokenizer.model_max_length,
            return_tensors="pt",
        )
        
        image_features = torch.stack([torch.Tensor(feature_extractor(image)["pixel_values"][0]) for image in images])
#         import pdb; pdb.set_trace()
        return {"input_ids": tokenized_captions["input_ids"], "images": image_features}

collate_fn = CollateFn()
laion_ds = datasets.load_dataset("laion/laion400m", split="train", streaming=True)
laion_ds_batched = laion_ds.map(collate_fn, batched=True, batch_size=8, remove_columns=next(iter(laion_ds)).keys())
```

### Steps to reproduce the bug

A reproducible [kaggle kernel ](https://www.kaggle.com/sachin/laion-hf-dataset/edit) can be found here.

### Expected behavior

Would expect `next(iter(laion_ds_batched))` to produce two tensors of shape `(batch_size, 77)` and `batch_size, image_shape`.

### Environment info

datasets==2.12.0
python==3.10
Hi ! in your code batching is **only used within** `map`, to process examples in batch. The dataset itself however is not batched and returns elements one by one.

To iterate on batches, you can do
```python
for batch in dataset.iter(batch_size=8):
    ...
```