cuda check script

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)