INNER CODE UNIT · Python

main

Shilin-LU/MACE · training.py:11

def main(conf):
    
    device = 'cuda' if torch.cuda.is_available() else 'cpu'

    # stage 1 & 2 (CFR and LoRA training)
    cfr_lora_training(conf.MACE)

    # stage 3 (Multi-LoRA fusion)
    multi_lora_fusion(conf.MACE)

    # test the erased model
    if conf.MACE.test_erased_model:
        inference(OmegaConf.create({
            "pretrained_model_name_or_path": conf.MACE.final_save_path,
            "multi_concept": conf.MACE.multi_concept,
            "generate_training_data": False,
            "device": device,
            "steps": 50,

View source record →