diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index 7abd49ad22..7048817559 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -383,8 +383,8 @@ def forward(
             # CUDA autocast runs RMSNorm in fp32, so cast explicitly before
             # caching: https://github.com/pytorch/pytorch/blob/v2.13.0/aten/src/ATen/autocast_mode.h#L875
             return out, KVCacheEntry(
-                key=key.to(value.dtype),
-                value=value,
+                key=key.to(value.dtype).contiguous(),
+                value=value.contiguous(),
             )
         return out