add missing import (#2113)

mosaicml · Mar 31, 2023 · 8b80a08 · 8b80a08
1 parent 98a2699
commit 8b80a08
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 1 deletion.
diff --git a/tests/models/test_hf_model.py b/tests/models/test_hf_model.py
@@ -359,6 +359,8 @@ def test_hf_loading_load_save_paths(checkpoint_upload_path: Optional[str], local
 def test_hf_loading_sentencepiece_tokenizer(modify_tokenizer: bool, tmp_path: Path, tiny_t5_model):
     transformers = pytest.importorskip('transformers')
 
+    from composer.models import HuggingFaceModel
+
     t0_pp_tokenizer = transformers.AutoTokenizer.from_pretrained('bigscience/T0pp')
 
     if modify_tokenizer:

diff --git a/tests/trainer/test_checkpoint.py b/tests/trainer/test_checkpoint.py
@@ -526,7 +526,7 @@ def get_trainer(self, **kwargs):
         model = SimpleModel()
         optimizer = torch.optim.Adam(model.parameters())
 
-        train_dataset = RandomClassificationDataset(size=25)
+        train_dataset = RandomClassificationDataset(size=20)
         eval_dataset = RandomClassificationDataset(size=10)
         train_batch_size = 2