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]],