Running alphafold on masks

How can I run alphafold on `seq_mask`, `msa_mask` and `msa_row_mask`? In alphafold training, msa are masked *randomly* but I am seeking to implement custom masking.

Here is what I have looked into.

In `processed_feature_dict['msa_mask']` under `run_alphafold.py`, apparently all msa are masked (with value 0) except the first row, presumably the target sequence. Manually masking `target_feat` and `msa_feat` with `seq_mask`, `msa_mask` and `aatype` does not seem to be effective. I.e. after masking the first 5 residues on a 12-residue-long sequence, alphafold recovers "native residues" with 0.99 chance. So I must be missing something here.

Thank you for any suggestions. If needed, am happy to provide more details on what I have tried.