diff --git a/walrus/models/spatial_blocks/full_attention.py b/walrus/models/spatial_blocks/full_attention.py
index 3d65040..6c9ddc7 100755
--- a/walrus/models/spatial_blocks/full_attention.py
+++ b/walrus/models/spatial_blocks/full_attention.py
@@ -113,7 +113,7 @@ def forward(self, x, bcs, return_att=False):
         input = x.clone()
         x = self.norm1(x)
 
-        fused_ff_qkv = rearrange(x, "b c h w d -> b h w d c")
+        fused_ff_qkv = rearrange(x, "b c h w d -> b h w d c").contiguous()
         ff, q, k, v = self.fused_ff_qkv(fused_ff_qkv).split(self.fused_dims, dim=-1)
 
         # Split into heads and process q, k