feat: raise initializer error on unavailable data

2022-01-31 12:27:48 +01:00 · 2022-01-31 12:27:48 +01:00 · 62df3c0457
commit 62df3c0457
parent 07a2d6caaa
1 changed files with 2 additions and 0 deletions
--- a/prototorch/core/initializers.py
+++ b/prototorch/core/initializers.py
@ -217,6 +217,8 @@ class AbstractStratifiedCompInitializer(AbstractClassAwareCompInitializer):
        components = torch.tensor([])
        for k, v in distribution.items():
            stratified_data = self.data[self.targets == k]
+            if len(stratified_data) == 0:
+              raise ValueError(f"No data available for class {k}.")
            initializer = self.subinit_type(
                stratified_data,
                noise=self.noise,