diff --git a/sdm/nn/attention.py b/sdm/nn/attention.py
index 199d8376f..e8a380392 100644
--- a/sdm/nn/attention.py
+++ b/sdm/nn/attention.py
@@ -171,7 +171,7 @@ def forward(
         key = key.transpose(-3, -2)  # [B, Hkv, KV, C]
         value = value.transpose(-3, -2)  # [B, Hkv, KV, C]
 
-        # CUDA attention kernels accept at most 65,535 batch entries per call:
+        # CUDA attention kernels accept at most 65,535 batch entries per call.
         outs = [
             F.scaled_dot_product_attention(
                 query=query[i : i + 65_535],
@@ -185,7 +185,11 @@ def forward(
                 is_causal=is_causal,
                 scale=self.scale,
             ).transpose(-3, -2)  # [B, Q, Hq, C]
-            for i in range(0, query.size(0), 65_535)
+            for i in (
+                (0,)
+                if query.size(0) <= 65_535
+                else range(0, query.size(0), 65_535)
+            )
         ]
         out = outs[0] if len(outs) == 1 else torch.cat(outs)