[feat] adding multiclass testing in shapes

uhlmanngroup · Aug 14, 2024 · d856741 · d856741
1 parent f994f31
commit d856741
Showing 1 changed file with 7 additions and 2 deletions.
diff --git 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,
     )