Container for multiple linear classifiers.
| 171 | |
| 172 | |
| 173 | class AllClassifiers(nn.Module): |
| 174 | """Container for multiple linear classifiers.""" |
| 175 | |
| 176 | def __init__(self, classifiers_dict: Dict[str, nn.Module]): |
| 177 | super().__init__() |
| 178 | self.classifiers_dict = nn.ModuleDict() |
| 179 | self.classifiers_dict.update(classifiers_dict) |
| 180 | |
| 181 | def forward(self, inputs) -> Dict[str, torch.Tensor]: |
| 182 | return {k: v.forward(inputs) for k, v in self.classifiers_dict.items()} |
| 183 | |
| 184 | def __len__(self): |
| 185 | return len(self.classifiers_dict) |
| 186 | |
| 187 | |
| 188 | # ============================================================================ |
no outgoing calls