Thanks. Do you know how is it applied on a dimension? I have `inputs_embeds`

in batch (bs, tokens, dim). I guess I should apply it on tokens, right? Before, I had the following function:

```
avg_input_embeds, _ = torch.max(inputs_embeds, 1)
```

I expect it match with a MaxPool with size 1:

```
pool = torch.nn.AdaptiveMaxPool1d(1)
avg_input_embeds2 = pool(inputs_embeds.permute(0, 2,1)).permute(0,2,1)
```

Yet, not equal, however I guess I need some permute or something

Update: Yes, they are equal with unsqueeze

```
pool(inputs_embeds.permute(0,2,1)).permute(0,2,1) == avg_input_embeds.unsqueeze(1)
tensor([[[True, True, True, ..., True, True, True]],
[[True, True, True, ..., True, True, True]],
[[True, True, True, ..., True, True, True]],
```