Update aggregator.py

openvinotoolkit · Sep 27, 2024 · 4a88a47 · 4a88a47
1 parent c93676d
commit 4a88a47
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/nncf/common/tensor_statistics/aggregator.py b/nncf/common/tensor_statistics/aggregator.py
@@ -31,6 +31,8 @@
     "Calibration dataset must not be empty. Please provide calibration dataset with at least one sample."
 )
 
+DATASET_SIZE_ERROR = "The dataset size ({}) is smaller than the subset size ({})."
+
 
 class StatisticsAggregator(ABC):
     """
@@ -82,6 +84,11 @@ def collect_statistics(self, model: TModel, graph: NNCFGraph) -> None:
             empty_statistics = False
         if empty_statistics:
             raise nncf.ValidationError(EMPTY_DATASET_ERROR)
+
+        if len(self.dataset) < subset_size:
+            raise nncf.UnsupportedDatasetError(  
+                DATASET_SIZE_ERROR.format(len(self.dataset), subset_size)
+            )    
 
     def register_statistic_points(self, statistic_points: StatisticPointsContainer) -> None:
         """