precommit

mosaicml · Oct 25, 2024 · 29cb815 · 29cb815
1 parent 8ea0d99
commit 29cb815
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 2 deletions.
diff --git a/llmfoundry/models/llm_embed/modeling_llm_embed.py b/llmfoundry/models/llm_embed/modeling_llm_embed.py
@@ -28,8 +28,8 @@
 from transformers import PreTrainedTokenizerBase
 from transformers.modeling_outputs import CausalLMOutputWithPast
 
-from llmfoundry.models.mpt.configuration_mpt import MPTConfig
 from llmfoundry.models.hf.hf_causal_lm import ComposerHFCausalLM
+from llmfoundry.models.mpt.configuration_mpt import MPTConfig
 from llmfoundry.models.mpt.modeling_mpt import MPTForCausalLM
 from llmfoundry.models.utils.config_moe_args import create_set_process_group
 

diff --git a/tests/models/llm_embed/test_embedding_finetune.py b/tests/models/llm_embed/test_embedding_finetune.py
@@ -179,7 +179,7 @@ def test_construct_model_distributed(
          patch('torch.distributed.get_rank', return_value=0), \
          patch('torch.distributed.barrier'), \
          patch('transformers.AutoModel.from_pretrained', return_value=mock_auto_model), \
-         patch('runtime_private_plugins.models.llm_embed.FinetuneEmbeddingModel.construct_model', return_value=mock_auto_model):
+         patch('llmfoundry.models.llm_embed.FinetuneEmbeddingModel.construct_model', return_value=mock_auto_model):
         model_instance: FinetuneEmbeddingModel = FinetuneEmbeddingModel(
             tokenizer=mock_tokenizer,
             pretrained_model_name_or_path='bert-base-uncased',