The flash_attention_2 option is not included in the args: https://github.com/cumc-dbmi/cehrbert/blob/fb5d8ab0df60306f9fce72b9d49c2b9463d73bb8/src/cehrbert/runners/hf_runner_argument_dataclass.py#L320