# The datasets.map() method doesn't keep tensor format from \`tokenizer\`

**URL:** <https://discuss.huggingface.co/t/the-datasets-map-method-doesnt-keep-tensor-format-from-tokenizer/25482>\
**Category:** 🤗Datasets\
**Created:** [November 3, 2022, 2:38pm UTC](https://discuss.huggingface.co/t/the-datasets-map-method-doesnt-keep-tensor-format-from-tokenizer/25482 "2022-11-03T14:38:22Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![PhillBoy](https://avatars.discourse-cdn.com/v4/letter/p/e274bd/32.png) [@PhillBoy](https://discuss.huggingface.co/u/PhillBoy)\
**Post date:** [November 3, 2022, 2:38pm UTC](https://discuss.huggingface.co/t/the-datasets-map-method-doesnt-keep-tensor-format-from-tokenizer/25482/1 "2022-11-03T14:38:22Z")

</div>

The `map()` method from a dataset does not retain the tensor that is selected in the `return_tensor` argument.

Code:

```auto
from transformers import AutoTokenizer
from datasets import Dataset
data = {
    "text":[
        "This is a test"
    ]
}
dataset = Dataset.from_dict(data)

model_name = 'roberta-large-mnli'
tokenizer = AutoTokenizer.from_pretrained(model_name,problem_type="multi_label_classification")
def tokenize_function(examples):
    return tokenizer(examples["text"], truncation=True,return_tensors="pt")
tokenized_datasets = dataset.map(tokenize_function, remove_columns=["text"],batched=True)
tokenized_datasets[0]

```

Output:

```auto
{'input_ids': [0, 713, 16, 10, 1296, 2], 'attention_mask': [1, 1, 1, 1, 1, 1]}

```

However, just using text input into the tokenizer does.

Code:

```auto
from transformers import AutoTokenizer
model_name = 'roberta-large-mnli'
tokenizer = AutoTokenizer.from_pretrained(model_name,problem_type="multi_label_classification")
tokenizer("This is a test", truncation=True,return_tensors="pt")

```

Output:

```auto
{'input_ids': tensor([[0, 713, 16, 10, 1296, 2]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}

```

Why is this the case?

---

<div class="post-metadata">

**Author:** ![lhoestq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lhoestq/32/52888_2.png) [@lhoestq](https://discuss.huggingface.co/u/lhoestq)\
**Post date:** [November 4, 2022, 2:58pm UTC](https://discuss.huggingface.co/t/the-datasets-map-method-doesnt-keep-tensor-format-from-tokenizer/25482/2 "2022-11-04T14:58:14Z")

</div>

The `Dataset` object stores the data in an Apache Arrow table - a super efficient columnar data format.  
Therefore your torch tensors are converted to arrow when written to disk.

To get tensors back you can enable set the output format of the dataset to be “torch”:

```python
tokenized_datasets = tokenized_datasets.with_format("torch")
tokenized_datasets[0]
# {'input_ids': tensor([[0, 713, 16, 10, 1296, 2]]), 'attention_mask': tensor([[1, 1, 1, 1, 1, 1]])}

```