Title: `map()` fails midway due to format incompatibility 
State: closed
Labels: description: Something isn't working
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

Comments:
Cc: @lhoestq 

You can end up with a list instead of a tensor if all the tensors inside the list can't be stacked together - can you make sure all your inputs are tensors with the same shape ?

Is there an easy way to ensure it?

You can make sure your `tokenize` function always return tensors of the same shape

I modified my `tokenize()` function to be like so:

```py
def tokenize(batch):
    return tokenizer(batch["text"], padding="longest")
```

so that the padding always happens w.r.t to the length of the longest sequence in a batch. The issue still persists. Is there any other way? 

tbh I though your first implementation was fine
```python
def tokenize(batch):
    return tokenizer(batch["text"], padding=True, truncation=True)
```

Maybe you can try to see what the erroring data looks like by adding a try/except in `get_test_accuracy` ?

This is what I got. 

For the non-erroring data, it looks like (without the labels):

```
tensor([[  101, 10047,  3110,  ...,     0,     0,     0],
        [  101,  1045,  2514,  ...,     0,     0,     0],
        [  101,  1045,  2514,  ...,     0,     0,     0],
        ...,
        [  101,  1045,  2005,  ...,     0,     0,     0],
        [  101,  1045,  2572,  ...,     0,     0,     0],
        [  101, 10047,  7481,  ...,     0,     0,     0]]) 128
tensor([[1, 1, 1,  ..., 0, 0, 0],
        [1, 1, 1,  ..., 0, 0, 0],
        [1, 1, 1,  ..., 0, 0, 0],
        ...,
        [1, 1, 1,  ..., 0, 0, 0],
        [1, 1, 1,  ..., 0, 0, 0],
        [1, 1, 1,  ..., 0, 0, 0]]) 128
```

For the erroring part:

```
[tensor([ 101, 1045, 2064, 2102, 2393, 3110, 2066, 2242, 6355, 3047, 2004, 2574,
        2004, 1996, 8629, 2357, 2125, 4299, 1045, 2071, 2424, 2009, 2006, 7858,
         102,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,    0,
           0,    0,    0,    0,    0,    0]), tensor([  101, 10047,  5458,  1997,  3110, 11654,  1998, 11055,   102,     0,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0]), tensor([  101,  1045,  2074,  2064,  2102,  6073,  1996,  3110,  2008,  2026,
        14982,  2000,  5587,  2203, 16650, 29563,  2030,  2569,  4506,  2052,
         2191,  1037,  2738, 11552,  2208, 17044, 14540,  2100,  3375,   102,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0,     0,     0,     0,     0,
            0,     0,     0,     0,     0,     0]),
...

[tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
        1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]), tensor([1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0,
        0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]),
...
```

I also tried investigating the shapes of the individual entries within a `batch` without the labels:

```py
def get_test_accuracy(model):
    def fn(batch):        
        try:
            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()}
        except:
            for k in batch:
                if k != "label":
                    for  i in range(len(batch[k])):
                        print(batch[k][i].shape)
    return fn
```

They are:

```
...
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([66])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
torch.Size([69])
```

There are differing shapes. I understand if I set `batch_size=None` in `emotions_encoded = emotions.map(tokenize, batched=True)` the problem should be fixed as the whole dataset would be treated as a single batch. But is there a way to do that in batches? 

If you use the same batch_size for your two maps, you should get the exact same batches - therefore all containing the same shapes

Oh I see. Thanks. Closing this issue.