import torch from models.ram import RAM device = 'cuda' if torch.cuda.is_available() else 'cpu' def get_model(): """ Load the model. :param str model_name: name of the model :param str device: device :param bool grayscale: if True, the model is trained on grayscale images :param bool train: if True, the model is trained :return: model """ model = RAM() return model