Tensor operation not functioning in dataset mapping
{'login': 'xinghaow99', 'id': 50691954, 'node_id': 'MDQ6VXNlcjUwNjkxOTU0', 'avatar_url': 'https://avatars.githubusercontent.com/u/50691954?v=4', 'gravatar_id': '', 'url': 'https://api.github.com/users/xinghaow99', 'html_url': 'https://github.com/xinghaow99', 'followers_url': 'https://api.github.com/users/xinghaow99/followers', 'following_url': 'https://api.github.com/users/xinghaow99/following{/other_user}', 'gists_url': 'https://api.github.com/users/xinghaow99/gists{/gist_id}', 'starred_url': 'https://api.github.com/users/xinghaow99/starred{/owner}{/repo}', 'subscriptions_url': 'https://api.github.com/users/xinghaow99/subscriptions', 'organizations_url': 'https://api.github.com/users/xinghaow99/orgs', 'repos_url': 'https://api.github.com/users/xinghaow99/repos', 'events_url': 'https://api.github.com/users/xinghaow99/events{/privacy}', 'received_events_url': 'https://api.github.com/users/xinghaow99/received_events', 'type': 'User', 'site_admin': False}
[{'id': 1935892857, 'node_id': 'MDU6TGFiZWwxOTM1ODkyODU3', 'url': 'https://api.github.com/repos/huggingface/datasets/labels/bug', 'name': 'bug', 'color': 'd73a4a', 'default': True, 'description': "Something isn't working"}]
[]
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)` . > 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. 
## 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

{'url': 'https://api.github.com/repos/huggingface/datasets/issues/5133/reactions', 'total_count': 0, '+1': 0, '-1': 0, 'laugh': 0, 'hooray': 0, 'confused': 0, 'heart': 0, 'rocket': 0, 'eyes': 0}