# Feature extraction output

**URL:** <https://discuss.huggingface.co/t/feature-extraction-output/15637>\
**Category:** Beginners\
**Created:** [March 12, 2022, 8:48pm UTC](https://discuss.huggingface.co/t/feature-extraction-output/15637 "2022-03-12T20:48:42Z")\
**Posts on this page:** 1\
**Page:** 1

<div class="post-metadata">

**Author:** ![virilo](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/virilo/32/3673_2.png) [@virilo](https://discuss.huggingface.co/u/virilo)\
**Post date:** [March 12, 2022, 8:48pm UTC](https://discuss.huggingface.co/t/feature-extraction-output/15637/1 "2022-03-12T20:48:42Z")

</div>

I’m feeding a model (AutoModel.from\_pretrained “distilbert-base-uncased”) with a batch of 64 samples, and its returns me a transformers.trainer\_utils.PredictionOutput object.

How could I extract the embeddings for these 64 rows?

I’m using:

```auto
with torch.no_grad():

    for chunk in tqdm(np.array_split(test_df, INFERENCE_BATCH_SIZE)):
        test_dataset=Dataset.from_pandas(chunk.loc[:,[COL_TEXT]]) 
        tokenized_test_dataset = test_dataset.map(tokenize_function, batched=True)
        tokenized_test_dataset=tokenized_test_dataset.remove_columns([x for x in tokenized_test_dataset.features.keys() if x not in ['input_ids', 'attention_mask']])
        
        test_result = trainer.predict(test_dataset=tokenized_test_dataset)
        ...

```

test\_result.\_\_class\_\_

> transformers.trainer\_utils.PredictionOutput

test\_result.predictions.shape

> (64, 133, 768)

I expected a output size of (64,768) having a context vector for each imput rows.

What this second dimension with a length of 133?

Are these the input multiplied by the attention layers?

Should I simply average them?

Thanks in advance.

Cheers,