From d856741c4cb692a7728bbd326b958ce6022b92b7 Mon Sep 17 00:00:00 2001 From: Craig Russell Date: Wed, 14 Aug 2024 11:39:07 +0100 Subject: [PATCH] [feat] adding multiclass testing in shapes --- bioimage_embed/shapes/tests/test_lightning.py | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/bioimage_embed/shapes/tests/test_lightning.py b/bioimage_embed/shapes/tests/test_lightning.py index d7fa1dc2..2e9f58b0 100644 --- a/bioimage_embed/shapes/tests/test_lightning.py +++ b/bioimage_embed/shapes/tests/test_lightning.py @@ -17,12 +17,17 @@ def transform(): ) +@pytest.fixture(params=[1, 2, 16]) +def classes(request): + return request.param + + @pytest.fixture -def dataset(transform): +def dataset(transform, classes): return FakeData( size=64, image_size=(3, 224, 224), - num_classes=2, + num_classes=classes, transform=transform, )