def METHOD_NAME(self, src_tokens, src_lengths, return_all_hiddens=False):
    self.attn_head_selector.head_select(self.task_ids)
    return super().METHOD_NAME(src_tokens, src_lengths, return_all_hiddens)