Hi again @lewtun,
Thanks for your second answer. I’m looking forward to trying it. Probably this weekend because today was too tight for me to do it.
I looked at the code concerning this swapping of dimension in
        hidden = self.DistilBERT(ids, attention_mask=mask, return_dict=False)
        hidden = hidden[0]
        hidden = hidden.permute(0, 2, 1)
        hidden = self.maxpool(hidden)

and the reason is that the DistilBERT layer outputs a tensor of dimension (n_batch, input_sequence_length, 768) where 768 is just the size of individual token embeddings as DistilBERT outputs. As the nn.MaxPool1d layer takes the max over the last dimension, I swapped the dimension because I want the max pooling to be done on the “temporal” axis so in other words take the max over the embeddings obtained for the different input tokens.
Concerning this averaging of unmasked hidden states, I’m a bit confused. I know what masking is, at least the general idea, but I don’t understand why we wouldn’t want to have the embeddings of the (randomly) masked tokens into our average.
Also, is averaging a more common practice than taking the max for these kinds of models (the transformer based ones)?
Anyway, thanks a lot for your support. It was really great. I’ll try all this and come back to you. Sorry for the delay.