diff --git a/docker/check_cuda_torch.py b/docker/check_cuda_torch.py new file mode 100644 index 0000000..af47470 --- /dev/null +++ b/docker/check_cuda_torch.py @@ -0,0 +1,25 @@ +''' +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) \ No newline at end of file