From fa246d52318dadb7b6ca429d0f674a7375abe9ca Mon Sep 17 00:00:00 2001 From: Fynn Schmitt-Ulms Date: Wed, 1 Apr 2026 10:29:33 -0400 Subject: [PATCH] Fix shape comment in extract_hidden_states example (#38723) Signed-off-by: Fynn Schmitt-Ulms --- examples/offline_inference/extract_hidden_states.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/examples/offline_inference/extract_hidden_states.py b/examples/offline_inference/extract_hidden_states.py index 61299101c..551f13761 100644 --- a/examples/offline_inference/extract_hidden_states.py +++ b/examples/offline_inference/extract_hidden_states.py @@ -54,5 +54,5 @@ with tempfile.TemporaryDirectory() as tmpdirname: print("Extracted token ids:", token_ids) # Matches prompt token ids print( "Extracted hidden states shape:", hidden_states.shape - ) # [num_hidden_layers, prompt len, hidden size] + ) # [prompt len, num_hidden_layers, hidden size] print("Extracted hidden states:", hidden_states)