Jan Kieseler
26 lines
876 B
Python
Executable File
26 lines
876 B
Python
Executable File
#!/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) |