    def __call__(self, module, module_in, module_out):
        """
        Called as forward hook.
        """
        # embed subword token IDs and mean pool to get
        # embedding of original word
        ids1_embeddings = []
        for i in self.ids1:
            i = i.to(module.weight.device)
            ids1_embeddings.append(
                torch.mean(module.forward(i), dim=0, keepdim=True)
            )  # forward() does not trigger hooks, thereby avoiding infinite recursion
        ids2_embeddings = []
        for i in self.ids2:
            i = i.to(module.weight.device)
            ids2_embeddings.append(torch.mean(module.forward(i), dim=0, keepdim=True))
        ids1_embeddings = torch.cat(ids1_embeddings)
        ids2_embeddings = torch.cat(ids2_embeddings)

        module_out_size = module_out.size()
        # flatten tensor except for last dimension
        module_out = module_out.flatten(end_dim=-2)
        # only return bias-mitigated evaluation embeddings
        module_out = self.mitigator(
            module_out,
            self.predetermined_bias_direction.to(module_out.device),
            ids1_embeddings.to(module_out.device),
            ids2_embeddings.to(module_out.device),
        )[: module_out.size(0)]
        return module_out.reshape(module_out_size)
Called as forward hook.
allennlp/fairness/bias_mitigator_wrappers.py