feat: raise initializer error on unavailable data

This commit is contained in:
Jensun Ravichandran 2022-01-31 12:27:48 +01:00
parent 07a2d6caaa
commit 62df3c0457
No known key found for this signature in database
GPG Key ID: 7612C0CAB643D921

View File

@ -217,6 +217,8 @@ class AbstractStratifiedCompInitializer(AbstractClassAwareCompInitializer):
components = torch.tensor([]) components = torch.tensor([])
for k, v in distribution.items(): for k, v in distribution.items():
stratified_data = self.data[self.targets == k] 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( initializer = self.subinit_type(
stratified_data, stratified_data,
noise=self.noise, noise=self.noise,