model_factory.py

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