nshtrainer 0.34.1__tar.gz → 0.34.2__tar.gz

{nshtrainer-0.34.1 → nshtrainer-0.34.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: nshtrainer
-Version: 0.34.1
+Version: 0.34.2
 Summary: 
 Author: Nima Shoghi
 Author-email: nimashoghi@gmail.com

{nshtrainer-0.34.1 → nshtrainer-0.34.2}/pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "nshtrainer"
-version = "0.34.1"
+version = "0.34.2"
 description = ""
 authors = ["Nima Shoghi <nimashoghi@gmail.com>"]
 readme = "README.md"

{nshtrainer-0.34.1 → nshtrainer-0.34.2}/src/nshtrainer/data/balanced_batch_sampler.py

@@ -81,6 +81,22 @@ class BalancedBatchSampler(BatchSampler):
     ):
         super().__init__(sampler, batch_size, drop_last=drop_last)
 
+        # Validate the dataset
+        dataset = self._unwrap_dataset(self.distributed_sampler.dataset)
+        # Dataset much either implement `data_sizes`, or we need to provide a custom
+        # implementation of the dataset sizes function.
+        if isinstance(dataset, DatasetWithSizes):
+            log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
+
+        elif self._data_sizes_fn is not None:
+            log.critical("BalancedBatchSampler: Using custom data_sizes_fn")
+        else:
+            raise ValueError(
+                "Dataset must implement the `data_sizes` method, "
+                "or a custom data_sizes_fn must be provided "
+                "to the BalancedBatchSampler."
+            )
+
         self._device = device
         self._data_sizes_fn = data_sizes_fn
 
@@ -97,7 +113,6 @@ class BalancedBatchSampler(BatchSampler):
         # Dataset much either implement `data_sizes`, or we need to provide a custom
         # implementation of the dataset sizes function.
         if isinstance(dataset, DatasetWithSizes):
-            log.critical(f"BalancedBatchSampler: Resolved dataset to {type(dataset)}")
             return dataset.data_sizes(indices)
 
         if (data_sizes_fn := self._data_sizes_fn) is not None: