From 7c5524074f1498d79f583c10d8eca79d4f726a54 Mon Sep 17 00:00:00 2001 From: SeanNaren Date: Wed, 22 Jun 2022 19:50:32 +0100 Subject: [PATCH] See if this fixes --- tests/task/vision/test_image_classification.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/task/vision/test_image_classification.py b/tests/task/vision/test_image_classification.py index f01c8246..62d2193d 100644 --- a/tests/task/vision/test_image_classification.py +++ b/tests/task/vision/test_image_classification.py @@ -17,7 +17,7 @@ def test_smoke_train(hf_cache_path): feature_extractor = AutoFeatureExtractor.from_pretrained(pretrained_model_name_or_path="nateraw/tiny-vit-random") dm = ImageClassificationDataModule( - cfg=ImageClassificationDataConfig(batch_size=2, dataset_name="beans"), + cfg=ImageClassificationDataConfig(batch_size=1, dataset_name="beans"), feature_extractor=feature_extractor, ) model = ImageClassificationTransformer(