Hey there everyone, I haven’t used this amazing API for a while. The code below is able to pool the [CLS] output embedding from DistilBERT?
distilbert = DistilBertModel.from_pretrained('distilbert-base-uncased', return_dict=True)

# more code ...

# input_ids & attention_mask from a batch
outputs = distilbert(input_ids, attention_mask=attention_mask)
seq_embeddings = outputs[0] # DistilBERT output embeddings 
cls_embedding = seq_embeddings[:,0,:] # <--- [CLS] output embedding

I’ve taken this idea from the following DPR code:
  

      github.com
  

  
    huggingface/transformers/blob/master/src/transformers/models/dpr/modeling_dpr.py#L204



      
      
          ) -> Union[BaseModelOutputWithPooling, Tuple[Tensor, ...]]:
              outputs = self.bert_model(
                  input_ids=input_ids,
                  attention_mask=attention_mask,
                  token_type_ids=token_type_ids,
                  inputs_embeds=inputs_embeds,
                  output_attentions=output_attentions,
                  output_hidden_states=output_hidden_states,
                  return_dict=return_dict,
              )
              sequence_output, pooled_output = outputs[:2]
              pooled_output = sequence_output[:, 0, :]
              if self.projection_dim > 0:
                  pooled_output = self.encode_proj(pooled_output)
          
          
    if not return_dict:
                  return (sequence_output, pooled_output) + outputs[2:]
          
          
    return BaseModelOutputWithPooling(
                  last_hidden_state=sequence_output,
                  pooler_output=pooled_output,