diff --git a/unsloth/models/cohere.py b/unsloth/models/cohere.py
index c33317ee02..25aa3c9960 100644
--- a/unsloth/models/cohere.py
+++ b/unsloth/models/cohere.py
@@ -296,7 +296,7 @@ def CohereAttention_fast_forward_inference(
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
             self.temp_O = torch.empty(
-                (1, bsz, hidden_size), dtype = dtype, device = "cuda:0"
+                (bsz, 1, hidden_size), dtype = dtype, device = "cuda:0"
             )
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]
diff --git a/unsloth/models/falcon_h1.py b/unsloth/models/falcon_h1.py
index 428f49d727..ac970d0b2f 100644
--- a/unsloth/models/falcon_h1.py
+++ b/unsloth/models/falcon_h1.py
@@ -265,7 +265,7 @@ def FalconH1Attention_fast_forward_inference(
 
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
-            self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+            self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]
 
diff --git a/unsloth/models/gemma2.py b/unsloth/models/gemma2.py
index 03e77f6504..b6e125d688 100644
--- a/unsloth/models/gemma2.py
+++ b/unsloth/models/gemma2.py
@@ -352,7 +352,7 @@ def Gemma2Attention_fast_forward_inference(
         )
         self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
         # Only for Gemma2
-        self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+        self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         self.attention = torch.empty(
             (bsz, n_heads, 1, KV_CACHE_INCREMENT + seq_len), dtype = dtype, device = device
         )
diff --git a/unsloth/models/granite.py b/unsloth/models/granite.py
index 168df90f4c..be753dad38 100644
--- a/unsloth/models/granite.py
+++ b/unsloth/models/granite.py
@@ -324,7 +324,7 @@ def GraniteAttention_fast_forward_inference(
         )
         self.RH_Q = torch.empty((bsz, n_heads, 1, head_dim), dtype = dtype, device = device)
         # Only for Gemma2
-        self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+        self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         self.attention = torch.empty(
             (bsz, n_heads, 1, KV_CACHE_INCREMENT + seq_len), dtype = dtype, device = device
         )
diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py
index a3fc9ab49a..23d29499a1 100644
--- a/unsloth/models/llama.py
+++ b/unsloth/models/llama.py
@@ -397,7 +397,7 @@ def LlamaAttention_fast_forward_inference(
 
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
-            self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+            self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]
 
diff --git a/unsloth/models/qwen3.py b/unsloth/models/qwen3.py
index ea06016d72..0773a68309 100644
--- a/unsloth/models/qwen3.py
+++ b/unsloth/models/qwen3.py
@@ -249,7 +249,7 @@ def Qwen3Attention_fast_forward_inference(
 
         # Mistral Nemo 12b has weird dimensions
         if attention_size != hidden_size:
-            self.temp_O = torch.empty((1, bsz, hidden_size), dtype = dtype, device = device)
+            self.temp_O = torch.empty((bsz, 1, hidden_size), dtype = dtype, device = device)
         else:
             self.temp_O = self.temp_QA[1][:, :, :hidden_size]