From 573bbb974629c45d650483ce53d0aa08074663af Mon Sep 17 00:00:00 2001
From: chenyu <chenyu@fastmail.com>
Date: Tue, 29 Apr 2025 15:28:38 -0400
Subject: [PATCH] Revert "remove TransformerBlock contiguous in llama (#10104)"
 (#10108)

This reverts commit b8d07dcc544e1228e101fdad5924453b4e7fb13a.
---
 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 6a62280e01..3a089facbc 100644
--- a/extra/models/llama.py
+++ b/extra/models/llama.py
@@ -107,7 +107,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))
+    return (h + self.feed_forward(self.ffn_norm(h))).contiguous()
 
 # standard openai sampling
 def sample(logits: Tensor, temp: float, k: int, p: float, af: float, ap: float):