diff --git a/comfy/text_encoders/llama.py b/comfy/text_encoders/llama.py
index 4415d6e9f..ddac547e0 100644
--- a/comfy/text_encoders/llama.py
+++ b/comfy/text_encoders/llama.py
@@ -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,