num_features=24

Lightning-AI · May 3, 2023 · aa6024a · aa6024a
1 parent e2dc026
commit aa6024a
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/tests/legacy/simple_classif_training.py b/tests/legacy/simple_classif_training.py
@@ -38,7 +38,7 @@ def main_train(dir_path, max_epochs: int = 20):
         deterministic=True,
     )
 
-    dm = ClassifDataModule()
+    dm = ClassifDataModule(num_features=24, length=800, num_classes=3)
     model = ClassificationModel(num_features=24, num_classes=3, lr=0.01)
     trainer.fit(model, datamodule=dm)
     res = trainer.test(model, datamodule=dm)