`map()` fails midway due to format incompatibility 
State: closed
Labels: name: bug

### Describe the bug

I am using the `emotion` dataset from Hub for sequence classification. After training the model, I am using it to generate predictions for all the entries present in the `validation` split of the dataset. 

```py
def get_test_accuracy(model):
    def fn(batch):
        inputs = {k:v.to(device) for k,v in batch.items() 
                  if k in tokenizer.model_input_names}
        with torch.no_grad():
            output = model(**inputs)
            pred_label = torch.argmax(output.logits, axis=-1)
        return {"predicted_label": pred_label.cpu().numpy()}
    return fn
```

This is how the `get_test_accuracy()` is being used:

```py

emotions = load_dataset("emotion")

def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)

emotions_encoded = emotions.map(tokenize, batched=True)
emotions_encoded.set_format("torch", 
                            columns=["input_ids", "attention_mask", "label"])

new_dataset = emotions_encoded["validation"].map(
    accuracy_fn, batched=True, batch_size=128
)
```

Complete code is available in the Colab Notebook provided below. 

The `map()` process fails midway giving:

```shell
AttributeError                            Traceback (most recent call last)
<ipython-input-8-ad24ac288eb4> in <module>
      2 
      3 new_dataset = emotions_encoded["validation"].map(
----> 4     accuracy_fn, batched=True, batch_size=128
      5 )

7 frames
/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in map(self, function, with_indices, with_rank, 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, desc)
   2588                 new_fingerprint=new_fingerprint,
   2589                 disable_tqdm=disable_tqdm,
-> 2590                 desc=desc,
   2591             )
   2592         else:

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    582             self: "Dataset" = kwargs.pop("self")
    583         # apply actual function
--> 584         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    585         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    586         for dataset in datasets:

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in wrapper(*args, **kwargs)
    549         }
    550         # apply actual function
--> 551         out: Union["Dataset", "DatasetDict"] = func(self, *args, **kwargs)
    552         datasets: List["Dataset"] = list(out.values()) if isinstance(out, dict) else [out]
    553         # re-apply format to the output

/usr/local/lib/python3.7/dist-packages/datasets/fingerprint.py in wrapper(*args, **kwargs)
    478             # Call actual function
    479 
--> 480             out = func(self, *args, **kwargs)
    481 
    482             # Update fingerprint of in-place transforms + update in-place history of transforms

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in _map_single(self, function, with_indices, with_rank, 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, disable_tqdm, desc, cache_only)
   2970                                 indices,
   2971                                 check_same_num_examples=len(input_dataset.list_indexes()) > 0,
-> 2972                                 offset=offset,
   2973                             )
   2974                         except NumExamplesMismatchError:

/usr/local/lib/python3.7/dist-packages/datasets/arrow_dataset.py in apply_function_on_filtered_inputs(inputs, indices, check_same_num_examples, offset)
   2850             if with_rank:
   2851                 additional_args += (rank,)
-> 2852             processed_inputs = function(*fn_args, *additional_args, **fn_kwargs)
   2853             if update_data is None:
   2854                 # Check if the function returns updated examples

<ipython-input-6-4e0d280426f6> in fn(batch)
      1 def get_test_accuracy(model):
      2     def fn(batch):
----> 3         inputs = {k:v.to(device) for k,v in batch.items() 
      4                   if k in tokenizer.model_input_names}
      5         with torch.no_grad():

<ipython-input-6-4e0d280426f6> in <dictcomp>(.0)
      2     def fn(batch):
      3         inputs = {k:v.to(device) for k,v in batch.items() 
----> 4                   if k in tokenizer.model_input_names}
      5         with torch.no_grad():
      6             output = model(**inputs)

AttributeError: 'list' object has no attribute 'to'
```

As you'd notice in the notebook, the process fails _midway_ and not at the beginning. 

Is this expected? 

### Steps to reproduce the bug

Colab Notebook: 

https://colab.research.google.com/gist/sayakpaul/d1570d537faf39040d02d77b1ed7de07/scratchpad.ipynb

### Expected behavior

The mapping process should complete as is. If you switch the `split` to `test` it works as expected. 

### Environment info

Colab