code: def test_initialize_pytorch_with_distributed(mock_torch):
      mock_torch.cuda.is_available.return_value = True
      mock_torch.cuda.device_count.return_value = 4
  
      with clean_params():
          initialize_pytorch(local_rank=1, local_size=4)
  
      mock_torch.cuda.set_device.assert_called_with(1)
      assert "CUDA_VISIBLE_DEVICES" not in os.environ
  
file_path: ludwig/tests/ludwig/utils/test_torch_utils.py
imports: import contextlib
  import os
  from typing import List
  from unittest.mock import patch
  import pytest
  import torch
  from ludwig.utils.torch_utils import _get_torch_init_params, _set_torch_init_params, initialize_pytorch, sequence_length_2D, sequence_length_3D
label_desc: Use image node_id 9 for calling a global function with example usage: test_initialize_pytorch_with_distributed(mock_torch) without return types