From 62df3c04573ae3d376469290843ea482d5f623df Mon Sep 17 00:00:00 2001
From: Jensun Ravichandran <jjensun@gmail.com>
Date: Mon, 31 Jan 2022 12:27:48 +0100
Subject: [PATCH] feat: raise initializer error on unavailable data

---
 prototorch/core/initializers.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/prototorch/core/initializers.py b/prototorch/core/initializers.py
index 909628a..a21fad8 100644
--- 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,