Method __init__

tools/test_linear_probing_hf.py:176–179

(self, classifiers_dict: Dict[str, nn.Module])

    """Container for multiple linear classifiers."""

    def __init__(self, classifiers_dict: Dict[str, nn.Module]):
        super().__init__()
        self.classifiers_dict = nn.ModuleDict()
        self.classifiers_dict.update(classifiers_dict)

    def forward(self, inputs) -> Dict[str, torch.Tensor]:
        return {k: v.forward(inputs) for k, v in self.classifiers_dict.items()}