Tensor operation not functioning in dataset mapping
COLLABORATOR
Hi! The Torch ops in your snippet are not equivalent to the NumPy ones, hence the difference. You can get the same behavior by replacing the line `feature = torch.mean(feature, dim=1)` with `feature = feature.squeeze().mean(1)` .
2022-10-18T18:16:06Z
https://github.com/huggingface/datasets/issues/5133#issuecomment-1282821416
https://api.github.com/repos/huggingface/datasets/issues/5133
IC_kwDODunzps5Mdk0o
https://api.github.com/repos/huggingface/datasets/issues/comments/1282821416/reactions
2022-10-18T18:16:06Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1282821416
https://avatars.githubusercontent.com/u/47462742?v=4
https://api.github.com/users/mariosasko/events{/privacy}
https://api.github.com/users/mariosasko/followers
https://api.github.com/users/mariosasko/following{/other_user}
https://api.github.com/users/mariosasko/gists{/gist_id}
https://github.com/mariosasko
mariosasko
MDQ6VXNlcjQ3NDYyNzQy
https://api.github.com/users/mariosasko/orgs
https://api.github.com/users/mariosasko/received_events
https://api.github.com/users/mariosasko/repos
https://api.github.com/users/mariosasko/starred{/owner}{/repo}
https://api.github.com/users/mariosasko/subscriptions
User
https://api.github.com/users/mariosasko
public
NONE
> Hi! The Torch ops in your snippet are not equivalent to the NumPy ones, hence the difference. You can get the same behavior by replacing the line `feature = torch.mean(feature, dim=1)` with `feature = feature.squeeze().mean(1)` .

Thank you. 
2022-10-19T04:15:44Z
https://github.com/huggingface/datasets/issues/5133#issuecomment-1283400106
https://api.github.com/repos/huggingface/datasets/issues/5133
IC_kwDODunzps5MfyGq
https://api.github.com/repos/huggingface/datasets/issues/comments/1283400106/reactions
2022-10-19T04:15:44Z
https://api.github.com/repos/huggingface/datasets/issues/comments/1283400106
https://avatars.githubusercontent.com/u/50691954?v=4
https://api.github.com/users/xinghaow99/events{/privacy}
https://api.github.com/users/xinghaow99/followers
https://api.github.com/users/xinghaow99/following{/other_user}
https://api.github.com/users/xinghaow99/gists{/gist_id}
https://github.com/xinghaow99
xinghaow99
MDQ6VXNlcjUwNjkxOTU0
https://api.github.com/users/xinghaow99/orgs
https://api.github.com/users/xinghaow99/received_events
https://api.github.com/users/xinghaow99/repos
https://api.github.com/users/xinghaow99/starred{/owner}{/repo}
https://api.github.com/users/xinghaow99/subscriptions
User
https://api.github.com/users/xinghaow99
public
## Describe the bug
I'm doing a torch.mean() operation in data preprocessing, and it's not working.

## Steps to reproduce the bug
```
from transformers import pipeline
import torch
import numpy as np
from datasets import load_dataset
device = 'cuda:0'
raw_dataset = load_dataset("glue", "sst2")
feature_extraction = pipeline('feature-extraction', 'bert-base-uncased', device=device)
def extracted_data(examples):
    # feature = torch.tensor(feature_extraction(examples['sentence'], batch_size=16), device=device)
    # feature = torch.mean(feature, dim=1)
    feature = np.asarray(feature_extraction(examples['sentence'], batch_size=16)).squeeze().mean(1)
    print(feature.shape)
    return {'feature': feature}

extracted_dataset = raw_dataset.map(extracted_data, batched=True, batch_size=16)
```

## Results
When running with torch.mean(), the shape printed out is [16, seq_len, 768], which is exactly the same before the operation. While numpy works just fine, which gives [16, 768].


## Environment info
- `datasets` version: 2.6.1
- Platform: Linux-4.4.0-142-generic-x86_64-with-glibc2.31
- Python version: 3.10.6
- PyArrow version: 9.0.0
- Pandas version: 1.5.0
