llama: fix non-local x path (#15580)

2026-09-21 13:38:08 -05:00 · 2026-08-13 12:21:26 -04:00
parent e535e59e13
commit af3d2153a7
1 changed files with 2 additions and 1 deletions
@@ -906,7 +906,8 @@ class Llama2_(nn.Module):
                past_kv.prepare(seq_len)

            def core():
-                _, current_kv = layer(
+                nonlocal x
+                x, current_kv = layer(
                    x=x,
                    attention_mask=mask,
                    freqs_cis=freqs_cis,