# Tokenization on dataset produces invalid pytorch tensor shape

**URL:** <https://discuss.huggingface.co/t/tokenization-on-dataset-produces-invalid-pytorch-tensor-shape/21009>\
**Category:** 🤗Datasets\
**Created:** [July 29, 2022, 6:58pm UTC](https://discuss.huggingface.co/t/tokenization-on-dataset-produces-invalid-pytorch-tensor-shape/21009 "2022-07-29T18:58:54Z")\
**Posts on this page:** 2\
**Page:** 1

<div class="post-metadata">

**Author:** ![cambonator](https://avatars.discourse-cdn.com/v4/letter/c/ed655f/32.png) [@cambonator](https://discuss.huggingface.co/u/cambonator)\
**Post date:** [July 29, 2022, 6:58pm UTC](https://discuss.huggingface.co/t/tokenization-on-dataset-produces-invalid-pytorch-tensor-shape/21009/1 "2022-07-29T18:58:54Z")

</div>

I’m attempting to tokenize text in a datasets object and then perform model inference with the tokenized data. I’m using very similar code to [this example in the datasets documentation](https://huggingface.co/docs/datasets/v2.4.0/en/use_dataset#tokenize-text).

The issue I’m running into is that the pytorch tensors for the tokens end up in the wrong shape. I have a reproducible example below where I compare tokenizing a single string compared to tokenizing the dataset with map(). These methods produce tensor shapes of `torch.Size([1, 23])` and `torch.Size([23])` respectively.

How can I adjust the dataset processing so that I get the expected tensor shape (`torch.Size([1, 23])`)?

```auto
from transformers import AutoTokenizer, AutoModelForSequenceClassification
from datasets import load_dataset
import time, os

model_arch="distilbert-base-uncased"

# import raw text data
raw_data = load_dataset("glue","cola",split='train').shuffle(seed=42).select(range(10))
text_field='sentence'

model=AutoModelForSequenceClassification.from_pretrained(model_arch)
tokenizer=AutoTokenizer.from_pretrained(model_arch)

# tokenize a single row
adhoc_tokenized_data=tokenizer(raw_data[0][text_field],return_tensors='pt')

# tokenized the entire raw text dataset
tokenized_data = raw_data.map(
    lambda x: tokenizer(x[text_field]),
    remove_columns=[text_field, 'idx']
)
tokenized_data=tokenized_data.rename_column('label', 'labels') # forward() expects 'labels'
tokenized_data.set_format('torch')
print(tokenized_data.format) 

# look at shape of tensors
print(adhoc_tokenized_data['input_ids'].shape) # torch.Size([1, 23])
print(tokenized_data[0]['input_ids'].shape) # torch.Size([23])

# Attempt to run inference on a single row of data --------------------------------------------

# inference works on the single row of tokenized data
print(model(**adhoc_tokenized_data))

# inference fails on the first row of the tokenized dataset object
# Error from forward(): IndexError: Dimension out of range (expected to be in range of [-1, 0], but got 1)
print(model(**tokenized_data[0]))

```

---

<div class="post-metadata">

**Author:** ![lhoestq](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lhoestq/32/64813_2.png) [@lhoestq](https://discuss.huggingface.co/u/lhoestq)\
**Post date:** [August 16, 2022, 9:53am UTC](https://discuss.huggingface.co/t/tokenization-on-dataset-produces-invalid-pytorch-tensor-shape/21009/2 "2022-08-16T09:53:36Z")

</div>

A model expects a batch as input, but you are passing `tokenized_data[0]` which is a single row.

Instead, you can pass a batch of one or several rows with `tokenized_data[:1]` or `tokenized_data[:n_rows]`:

```python
print(model(**tokenized_data[:1]))

```