Issue: `ArrowInvalid` occurs while running `Dataset.map()` function for DPRContext

It seems to fail the final batch ):

steps to reproduce:
```
from datasets import load_dataset
from elasticsearch import Elasticsearch
import torch
from transformers import file_utils, set_seed
from transformers import DPRContextEncoder, DPRContextEncoderTokenizerFast
MAX_SEQ_LENGTH = 256
ctx_encoder = DPRContextEncoder.from_pretrained("facebook/dpr-ctx_encoder-single-nq-base", cache_dir="../datasets/")
ctx_tokenizer = DPRContextEncoderTokenizerFast.from_pretrained(
    "facebook/dpr-ctx_encoder-single-nq-base", 
    cache_dir="..datasets/"
)

dataset = load_dataset('text', 
                data_files='data/raw/ARC_Corpus.txt',
                cache_dir='../datasets')

torch.set_grad_enabled(False)
ds_with_embeddings = dataset.map(
    lambda example: {
        'embeddings': ctx_encoder(
            **ctx_tokenizer(
                example["text"], 
                padding='max_length', 
                truncation=True, 
                max_length=MAX_SEQ_LENGTH,
                return_tensors="pt"
            )
        )[0][0].numpy(),
    },
    batched=True,
    load_from_cache_file=False,
    batch_size=1000
)
```
ARC Corpus can be obtained from [here](https://ai2-datasets.s3-us-west-2.amazonaws.com/arc/ARC-V1-Feb2018.zip)

And then the error:

```
---------------------------------------------------------------------------
ArrowInvalid                              Traceback (most recent call last)
<ipython-input-13-67d139bb2ed3> in <module>
     14     batched=True,
     15     load_from_cache_file=False,
---> 16     batch_size=1000
     17 )

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/datasets/dataset_dict.py in map(self, function, with_indices, input_columns, batched, batch_size, remove_columns, keep_in_memory, load_from_cache_file, cache_file_names, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc)
    301                     num_proc=num_proc,
    302                 )
--> 303                 for k, dataset in self.items()
    304             }
    305         )

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/datasets/dataset_dict.py in <dictcomp>(.0)
    301                     num_proc=num_proc,
    302                 )
--> 303                 for k, dataset in self.items()
    304             }
    305         )

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/datasets/arrow_dataset.py in map(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, num_proc, suffix_template, new_fingerprint)
   1257                 fn_kwargs=fn_kwargs,
   1258                 new_fingerprint=new_fingerprint,
-> 1259                 update_data=update_data,
   1260             )
   1261         else:

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    155         }
    156         # apply actual function
--> 157         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    158         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    159         # re-apply format to the output

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    161             # Call actual function
    162 
--> 163             out = func(self, *args, **kwargs)
    164 
    165             # Update fingerprint of in-place transforms + update in-place history of transforms

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, input_columns, batched, batch_size, drop_last_batch, remove_columns, keep_in_memory, load_from_cache_file, cache_file_name, writer_batch_size, features, disable_nullable, fn_kwargs, new_fingerprint, rank, offset, update_data)
   1526                     if update_data:
   1527                         batch = cast_to_python_objects(batch)
-> 1528                         writer.write_batch(batch)
   1529             if update_data:
   1530                 writer.finalize()  # close_stream=bool(buf_writer is None))  # We only close if we are writing in a file

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/datasets/arrow_writer.py in write_batch(self, batch_examples, writer_batch_size)
    276             typed_sequence = TypedSequence(batch_examples[col], type=col_type, try_type=col_try_type)
    277             typed_sequence_examples[col] = typed_sequence
--> 278         pa_table = pa.Table.from_pydict(typed_sequence_examples)
    279         self.write_table(pa_table)
    280 

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_pydict()

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.from_arrays()

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/pyarrow/table.pxi in pyarrow.lib.Table.validate()

~/.cache/pypoetry/virtualenvs/masters-utTTC0p8-py3.7/lib/python3.7/site-packages/pyarrow/error.pxi in pyarrow.lib.check_status()

ArrowInvalid: Column 1 named text expected length 768 but got length 1000
```

Comments:
Looks like the mapping function returns a dictionary with a 768-dim array in the `embeddings` field. Since the map is batched, we actually expect the `embeddings` field to be an array of shape (batch_size, 768) to have one embedding per example in the batch.

To fix that can you try to remove one of the `[0]` ? In my opinion you only need one of them, not two.

It makes sense :D

It seems to work! Thanks a lot :))

Closing the issue