    def __call__(self, module, module_in, module_out):
        """
        Called as forward hook.
        """
        module_out_size = module_out.size()
        # flatten tensor except for last dimension
        module_out = module_out.flatten(end_dim=-2)
        module_out = self.mitigator(
            module_out, self.predetermined_bias_direction.to(module_out.device)
        )
        return module_out.reshape(module_out_size)
Called as forward hook.
allennlp/fairness/bias_mitigator_wrappers.py