Log dataset sizes (#246)

src/metatrain/cli/train.py

@@ -308,6 +308,31 @@ def train_model(
             )
             validation_datasets.append(validation_dataset)
 
+    ###########################
+    # PRINT DATASET STATS #####
+    ###########################
+
+    for i, train_dataset in enumerate(train_datasets):
+        if len(train_datasets) == 1:
+            index = ""
+        else:
+            index = f" {i}"
+        logger.info(f"Training dataset{index} has size {len(train_dataset)}")
+
+    for i, validation_dataset in enumerate(validation_datasets):
+        if len(validation_datasets) == 1:
+            index = ""
+        else:
+            index = f" {i}"
+        logger.info(f"Validation dataset{index} has size {len(validation_dataset)}")
+
+    for i, test_dataset in enumerate(test_datasets):
+        if len(test_datasets) == 1:
+            index = ""
+        else:
+            index = f" {i}"
+        logger.info(f"Test dataset{index} has size {len(test_dataset)}")
+
     ###########################
     # SAVE EXPANDED OPTIONS ###
     ###########################

tests/cli/test_train_model.py

@@ -72,6 +72,9 @@ def test_train(capfd, monkeypatch, tmp_path, output):
     for logtext in [stdout_log, file_log]:
         assert "This log is also available"
         assert re.search(r"Random seed of this run is [1-9]\d*", logtext)
+        assert "Training dataset has size"
+        assert "Validation dataset has size"
+        assert "Test dataset has size"
         assert "[INFO]" in logtext
         assert "Epoch" in logtext
         assert "loss" in logtext