#!/usr/bin/env python
'''
little script to check if cuda is available in pytorch and working
organised as a check that raises an exception if cuda is not available, and data cannot be moved to the gpu
'''
import torch
import numpy as np
import os

def check_cuda():
    if not torch.cuda.is_available():
        raise Exception("CUDA is not available. Please check if you have installed the correct version of CUDA and the correct version of the NVIDIA driver.")
    
    try:
        a = torch.tensor([1,2,3])
        a = a.cuda()
        a = a*a
    except:
        raise Exception("Cannot move data to the GPU. Please check if you have installed the correct version of CUDA and the correct version of the NVIDIA driver.")
    
    print("CUDA is available and working correctly.")

if __name__ == '__main__':
    check_cuda()
    print("All checks passed.")
    os._exit(0)