Extracting First Hidden States
## ❓ Questions & Help

For example, right now in order to extract the first hidden states of DistilBert model,

tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
model = DistilBertModel.from_pretrained('distilbert-base-uncased', **output_hidden_states=True**)
input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
outputs = model(input_ids)
**first_hidden_states = outputs[1][0]  # The first hidden-state**

However, I can only extract the first layer hidden states after running through the entire model; all 7 layers (which is unnecessary if I only want the first).

Is there a way I could use the model to only extract the hidden states of the first layer and stop there? I am looking to further optimize the inference time of DistilBert.

Thank you :)

