src/transformers/models/esm/modeling_esmfold.py/EsmForProteinFolding/af2_idx_to_esm_idx
class EsmForProteinFolding:    def af2_idx_to_esm_idx(self, aa, mask):
        # avoid indexing on different devices
        if self.af2_to_esm.device != aa.device:
            self.af2_to_esm = self.af2_to_esm.to(aa.device)
        aa = (aa + 1).masked_fill(mask != 1, 0)
        return self.af2_to_esm[aa]