Support returning last hidden state (#2461)

ebsmothers · web-flow · commit b4d7fbb6a0ef · 2025-03-05T13:54:10.000-08:00
diff --git a/torchtune/modules/transformer.py b/torchtune/modules/transformer.py
@@ -653,6 +653,9 @@ def forward(
                 input_pos=input_pos,
             )
 
+        if len(self.layers) in self.output_hidden_states:
+            hidden.append(h)
+
         # shape: [b, seq_len, out_dim]
         output = self.unembed(h)
 

Original file line number	Diff line number	Diff line change
`@@ -653,6 +653,9 @@ def forward(`
`653`	`653`	`input_pos=input_pos,`
`654`	`654`	`)`
`655`	`655`
	`656`	`+ if len(self.layers) in self.output_hidden_states:`
	`657`	`+ hidden.append(h)`
	`658`	`+`
`656`	`659`	`# shape: [b, seq_len, out_dim]`
`657`	`660`	`output = self.unembed(h)`
`658`	`661`