trainer/trainer.py

from .utils import *
from configs.configs import root,config,mae_config

def main():
    try:
        decision=input("train mae or classifier? ")
        if decision=="mae":
            print(f"Training mae")
            trainer=MAETrainer(mae_config)
            trainer.train()
        if decision=="classifier":
            print(f"Training classifier")
            trainer=Trainer(config)
            trainer.train()
    except:
        import traceback
        traceback.print_exc()
        
if __name__=="__main__":main()