# TypeError when applying map after set\_format(type='torch')

**URL:** <https://discuss.huggingface.co/t/typeerror-when-applying-map-after-set-format-type-torch/23067>\
**Category:** 🤗Datasets\
**Created:** [September 13, 2022, 6:57am UTC](https://discuss.huggingface.co/t/typeerror-when-applying-map-after-set-format-type-torch/23067 "2022-09-13T06:57:32Z")\
**Posts on this page:** 4\
**Page:** 1

<div class="post-metadata">

**Author:** ![simonschoe](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/simonschoe/32/20951_2.png) [@simonschoe](https://discuss.huggingface.co/u/simonschoe)\
**Post date:** [September 13, 2022, 6:57am UTC](https://discuss.huggingface.co/t/typeerror-when-applying-map-after-set-format-type-torch/23067/1 "2022-09-13T06:57:32Z")

</div>

Hi there,

I am trying to run a simple `forward_pass` function on my dataset producing the following error:

```auto
dataset['train'].set_format(type='torch', columns=['input_ids', 'attention_mask', 'label'])

def forward_pass(batch):
    input_ids = torch.tensor(batch['input_ids']).to(device)
    attention_mask = torch.tensor(batch['attention_mask']).to(device)
    with torch.no_grad():
        batch['logits'] = model(input_ids, attention_mask)['logits'].cpu().numpy()
    return batch

dataset['train'].map(forward_pass, batched=True, batch_size=16)

```

```auto
TypeError: Provided `function` which is applied to all elements of table returns a `dict` of types [<class 'torch.Tensor'>, <class 'torch.Tensor'>, <class 'torch.Tensor'>, <class 'numpy.ndarray'>]. When using `batched=True`, make sure provided `function` returns a `dict` of types like `(<class 'list'>, <class 'numpy.ndarray'>)`.

```

The error does not occur when I convert to `numpy` instead of `torch`:

```auto
dataset['train'].set_format(type='numpy', columns=['input_ids', 'attention_mask', 'label'])

```

Why is that the case? I couldn’t quite wrap my head around why the `map` call doesn’t handle the tensor data but is fine with using the numpy arrays? Super grateful for insights on the inner workings of the employed function! 🙂

Best  
Simon

---

<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:** [September 13, 2022, 12:57pm UTC](https://discuss.huggingface.co/t/typeerror-when-applying-map-after-set-format-type-torch/23067/2 "2022-09-13T12:57:27Z")

</div>

Hi ! I think this is an issue in `datasets`, let me open a PR to fix this.

In the meantime you need to make sure every tensor in the returned `batch` dict should be lists or numpy arrays

---

<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:** [September 13, 2022, 1:16pm UTC](https://discuss.huggingface.co/t/typeerror-when-applying-map-after-set-format-type-torch/23067/3 "2022-09-13T13:16:55Z")

</div>

Pull request is open at [Fix map batched with torch output + test by lhoestq · Pull Request #4972 · huggingface/datasets · GitHub](https://github.com/huggingface/datasets/pull/4972)

---

<div class="post-metadata">

**Author:** ![simonschoe](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/simonschoe/32/20951_2.png) [@simonschoe](https://discuss.huggingface.co/u/simonschoe)\
**Post date:** [September 13, 2022, 2:33pm UTC](https://discuss.huggingface.co/t/typeerror-when-applying-map-after-set-format-type-torch/23067/4 "2022-09-13T14:33:26Z")

</div>

Great, thanks for the quick reply and PR. The workaround with converting back to `numpy` works just fine, however, the fix will make the whole workflow more convenient and robust IMO.