fix

Signed-off-by: Hemil Desai <[email protected]>

Author: hemildesai

tests/unit_tests/_transformers/test_auto_model.py

@@ -149,9 +149,6 @@ def test_from_pretrained_registry_downloads_checkpoint_files_rank0(self):
             patch.object(transformers.AutoModelForCausalLM, "from_pretrained") as mock_hf_loader,
             patch("nemo_automodel._transformers.auto_model._get_resolved_checkpoint_files") as mock_get_files,
             patch("nemo_automodel._transformers.auto_model.os.path.isdir", return_value=False),
-            patch("nemo_automodel._transformers.auto_model.dist.is_initialized", return_value=True),
-            patch("nemo_automodel._transformers.auto_model.dist.get_world_size", return_value=1),
-            patch("nemo_automodel._transformers.auto_model.dist.get_rank", return_value=0),
             patch("nemo_automodel.components.distributed.utils.FirstRankPerNode") as mock_barrier,
         ):
             # Prepare a fake config with architectures and commit hash
@@ -187,9 +184,6 @@ def test_from_pretrained_registry_downloads_when_dist_uninitialized(self):
             patch.object(transformers.AutoModelForCausalLM, "from_pretrained") as mock_hf_loader,
             patch("nemo_automodel._transformers.auto_model._get_resolved_checkpoint_files") as mock_get_files,
             patch("nemo_automodel._transformers.auto_model.os.path.isdir", return_value=False),
-            patch("nemo_automodel._transformers.auto_model.dist.is_initialized", return_value=False),
-            patch("nemo_automodel._transformers.auto_model.dist.get_world_size", return_value=1),
-            patch("nemo_automodel._transformers.auto_model.dist.barrier") as mock_barrier,
         ):
             # Prepare a fake config with architectures and commit hash
             cfg = Mock()
@@ -213,8 +207,6 @@ def test_from_pretrained_registry_downloads_when_dist_uninitialized(self):
             _, kwargs = mock_get_files.call_args
             assert kwargs["pretrained_model_name_or_path"] == "dummy/repo-id"
             assert kwargs["commit_hash"] == "commit456"
-            # No barrier when dist not initialized
-            mock_barrier.assert_not_called()
 
     def test_from_config_happy_path(self):
         """Test the basic from_config functionality works."""