code: def test_split_tensor_along_last_dim(self):
      for tensor_model_paralell_world_size in range(
          1, self.world_size + 1
      ):
          if self.world_size % tensor_model_paralell_world_size > 0:
              continue
          parallel_state.initialize_model_parallel(
              tensor_model_parallel_size_=tensor_model_paralell_world_size
          )
  
          device = "cpu"
          input_tensor = torch.randn((100, 100, 100), device=device)
          splits = utils.split_tensor_along_last_dim(input_tensor, 10)
          last_dim_shapes = torch.tensor(
              [int(split.size()[-1]) for split in splits]
          )
  
          self.assertTrue(
              torch.equal(
                  last_dim_shapes,
                  torch.full((10,), 10),
              ),
              msg=f"tensor_model_paralell_world_size: {tensor_model_paralell_world_size}",
          )
  
          parallel_state.destroy_model_parallel()
  
file_path: apex/tests/L0/run_transformer/test_transformer_utils.py
imports: import logging
  import torch
  from torch.testing._internal import common_utils
  from apex.transformer import parallel_state
  from apex.transformer.tensor_parallel import utils
  from apex.transformer.testing.distributed_test_base import NcclDistributedTestBase
label_desc: Use image node_id 1 for calling the TransformerUtilsTest obj's underlying member method code with example usage: obj.test_split_tensor_along_last_dim() without return types