diff --git a/diffsynth/models/minimax_h3_dit.py b/diffsynth/models/minimax_h3_dit.py
index 3ed30de65..30025ce86 100644
--- a/diffsynth/models/minimax_h3_dit.py
+++ b/diffsynth/models/minimax_h3_dit.py
@@ -137,9 +137,7 @@ def forward(self, x, *, rope_freqs, cu_seqlens, max_seqlen=None):
         total = x.shape[0]
         qkv = self.qkv_proj(x)
         qkv = qkv.view(total, self.num_heads, 3, self.head_dim)
-        q = qkv[:, :, 0, :]
-        k = qkv[:, :, 1, :]
-        v = qkv[:, :, 2, :]
+        q, k, v = qkv.unbind(dim=2)
         q = self.q_norm(q)
         k = self.k_norm(k)
         if rope_freqs is not None:
diff --git a/diffsynth/models/minimax_h3_dit_comfy.py b/diffsynth/models/minimax_h3_dit_comfy.py
index 50a6d1b7a..9bf5de8f7 100644
--- a/diffsynth/models/minimax_h3_dit_comfy.py
+++ b/diffsynth/models/minimax_h3_dit_comfy.py
@@ -18,9 +18,7 @@ def _comfy_attention_forward(self, x, *, rope_freqs, cu_seqlens, max_seqlen=None
     total = x.shape[0]
     qkv = self.qkv_proj(x)
     qkv = qkv.view(total, 3, self.num_heads, self.head_dim)
-    q = qkv[:, 0, :, :]
-    k = qkv[:, 1, :, :]
-    v = qkv[:, 2, :, :]
+    q, k, v = qkv.unbind(dim=1)
     q = self.q_norm(q)
     k = self.k_norm(k)
     if rope_freqs is not None: