diff --git a/tests/unit_tests/tensor_parallel/test_tensor_parallel_utils.py b/tests/unit_tests/tensor_parallel/test_tensor_parallel_utils.py
index 5df774e5fff..334f973ca3b 100644
--- a/tests/unit_tests/tensor_parallel/test_tensor_parallel_utils.py
+++ b/tests/unit_tests/tensor_parallel/test_tensor_parallel_utils.py
@@ -1,3 +1,5 @@
+# Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
 import torch
 
 import megatron.core.parallel_state as ps
@@ -9,8 +11,9 @@
 
 def test_split_tensor_along_last_dim():
     input_tensor = torch.rand((3, 4))
-    torch.equal(input_tensor[0:2, 0:2], util.split_tensor_along_last_dim(input_tensor, 2)[0])
-    torch.equal(input_tensor[2:, 2:], util.split_tensor_along_last_dim(input_tensor, 2)[1])
+    chunks = util.split_tensor_along_last_dim(input_tensor, 2)
+    assert torch.equal(chunks[0], input_tensor[:, 0:2])
+    assert torch.equal(chunks[1], input_tensor[:, 2:4])
 
 
 def test_split_tensor_into_1d_equal_chunks():