Allow empty validation set in fit.

neuronets · Aug 31, 2023 · 5c1738b · 5c1738b
1 parent 4106a2d
commit 5c1738b
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 6 deletions.
diff --git a/nobrainer/processing/segmentation.py b/nobrainer/processing/segmentation.py
@@ -89,8 +89,10 @@ def _compile():
             dataset_train.dataset,
             epochs=epochs,
             steps_per_epoch=dataset_train.get_steps_per_epoch(),
-            validation_data=dataset_validate.dataset,
-            validation_steps=dataset_validate.get_steps_per_epoch(),
+            validation_data=dataset_validate.dataset if dataset_validate
+            else None,
+            validation_steps=dataset_validate.get_steps_per_epoch()
+            if dataset_validate else None,
             callbacks=callbacks,
         )
 

diff --git a/nobrainer/tests/checkpoint_test.py b/nobrainer/tests/checkpoint_test.py
@@ -6,6 +6,7 @@
 from numpy.testing import assert_allclose
 import tensorflow as tf
 
+from nobrainer.dataset import Dataset
 from nobrainer.models import meshnet
 from nobrainer.processing.segmentation import Segmentation
 
@@ -15,10 +16,7 @@ def _get_toy_dataset():
     train = tf.data.Dataset.from_tensors(
         (np.random.rand(*data_shape), np.random.randint(0, 1, data_shape))
     )
-    train.scalar_labels = False
-    train.n_volumes = data_shape[0]
-    train.volume_shape = data_shape[1:4]
-    return train
+    return Dataset(train, data_shape[0], data_shape[1:4], 1)
 
 
 def _assert_model_weights_allclose(model1, model2):