Revert "Fix error when batch size is larger than the dataset" (#616)

This reverts commit 4b7c9a4.

Author: Luthaf

src/metatrain/experimental/nanopet/trainer.py

@@ -149,7 +149,7 @@ def train(
                     num_replicas=world_size,
                     rank=rank,
                     shuffle=True,
-                    drop_last=len(train_dataset) > self.hypers["batch_size"],
+                    drop_last=True,
                 )
                 for train_dataset in train_datasets
             ]
@@ -181,9 +181,7 @@ def train(
                     ),
                     drop_last=(
                         # the sampler takes care of this (if present)
-                        # check if batch size > train_dataset
-                        len(train_dataset) > self.hypers["batch_size"]
-                        and train_sampler is None
+                        train_sampler is None
                     ),
                     collate_fn=collate_fn,
                 )

src/metatrain/pet/trainer.py

@@ -167,7 +167,7 @@ def train(
                     num_replicas=world_size,
                     rank=rank,
                     shuffle=True,
-                    drop_last=len(train_dataset) > self.hypers["batch_size"],
+                    drop_last=True,
                 )
                 for train_dataset in train_datasets
             ]
@@ -199,9 +199,7 @@ def train(
                     ),
                     drop_last=(
                         # the sampler takes care of this (if present)
-                        # check if batch size > train_dataset
-                        len(train_dataset) > self.hypers["batch_size"]
-                        and train_sampler is None
+                        train_sampler is None
                     ),
                     collate_fn=collate_fn,
                 )

src/metatrain/soap_bpnn/trainer.py

@@ -148,7 +148,7 @@ def train(
                     num_replicas=world_size,
                     rank=rank,
                     shuffle=True,
-                    drop_last=len(train_dataset) > self.hypers["batch_size"],
+                    drop_last=True,
                 )
                 for train_dataset in train_datasets
             ]
@@ -180,9 +180,7 @@ def train(
                     ),
                     drop_last=(
                         # the sampler takes care of this (if present)
-                        # check if batch size > train_dataset
-                        len(train_dataset) > self.hypers["batch_size"]
-                        and train_sampler is None
+                        train_sampler is None
                     ),
                     collate_fn=collate_fn,
                 )

tests/cli/test_train_model.py

@@ -302,19 +302,6 @@ def test_empty_training_set(monkeypatch, tmp_path, options):
         train_model(options)
 
 
-def test_batch_size_smaller_training_set(monkeypatch, tmp_path, options):
-    """Test that training still runs for batch size > train_size."""
-    monkeypatch.chdir(tmp_path)
-
-    shutil.copy(DATASET_PATH_QM9, "qm9_reduced_100.xyz")
-
-    options["validation_set"] = 0.55
-    options["test_set"] = 0.4
-    options["architecture"]["training"]["batch_size"] = 1000
-
-    train_model(options)
-
-
 @pytest.mark.parametrize("split", [-0.1, 1.1])
 def test_wrong_test_split_size(split, monkeypatch, tmp_path, options):
     """Test that an error is raised if the test split has the wrong size"""