diff --git a/tests/unit/module_inject/test_fused_repartition.py b/tests/unit/module_inject/test_fused_repartition.py
index f971e4ebf65f..5bf54d1a207b 100644
--- a/tests/unit/module_inject/test_fused_repartition.py
+++ b/tests/unit/module_inject/test_fused_repartition.py
@@ -64,7 +64,7 @@ def test_gate_up_partition_covers_the_whole_weight(clean_tp_shard_globals):
 
     gate = torch.cat([shards[0][:3], shards[1][:2]], dim=0)
     up = torch.cat([shards[0][3:], shards[1][2:]], dim=0)
-    assert torch.equal(torch.cat([gate, up], dim=0), full_weight)
+    assert torch.equal(torch.cat([gate, up], dim=0), full_weight.to(gate.device))
 
 
 class _QWenAttention(torch.nn.Module):