models.py

import os
import sys
import torch


def load_models():
    '''
    Checks CUDA availability & loads models
    '''
    try:
        print(f"CUDA Available: {torch.cuda.is_available()}")
        print(f"CUDA Device: {torch.cuda.get_device_name(torch.cuda.current_device())}")
        
        # TODO: Load Models Code Here (Refactor Exisitng Models Loading)
        MODELS = {}
        return MODELS

    except KeyboardInterrupt:
        print('Interrupted')
        try:
            sys.exit(0)
        except SystemExit:
            os._exit(0)