Tensor operation not functioning in dataset mapping
## 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

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. 