trains/ATIO.py

"""
ATIO -- All Trains in One
"""
from .singleTask import *

__all__ = ['ATIO']

class ATIO():
    def __init__(self):
        self.TRAIN_MAP = {
            'DLF': DLF,
        }
    
    def getTrain(self, args):
        return self.TRAIN_MAP[args['model_name']](args)