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)