ola_vlm/train/old/probe_train_mem.py

from ola_vlm.train.probe_train import train

if __name__ == "__main__":
    try:
        train(attn_implementation="flash_attention_2")
    except:
        train(attn_implementation="eager")