From 474ee9daa5fc25aefa2c19df16a871daa1ff3821 Mon Sep 17 00:00:00 2001 From: George Hotz Date: Thu, 31 Jul 2025 15:06:19 -0700 Subject: [PATCH] hotfix: add contiguous_backward to llama --- extra/models/llama.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/extra/models/llama.py b/extra/models/llama.py index ecee88f64a..09907b964e 100644 --- a/extra/models/llama.py +++ b/extra/models/llama.py @@ -111,7 +111,7 @@ class TransformerBlock: def __call__(self, x:Tensor, start_pos:Union[Variable,int], freqs_cis:Tensor, mask:Optional[Tensor]): h = x + self.attention(self.attention_norm(x), start_pos, freqs_cis, mask) - return (h + self.feed_forward(self.ffn_norm(h))).contiguous() + return (h + self.feed_forward(self.ffn_norm(h))).contiguous().contiguous_backward() # standard openai sampling def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):