import torch

import models

#__all__={'make'}

models = {}

def register(name):
    def decorator(cls):
        models[name] = cls
        return cls
    return decorator

def make(device, name, **kwargs):
    if name is None:
        return None
    
    model = models[name](**kwargs)
    model.to(device)

    return model