-
Notifications
You must be signed in to change notification settings - Fork 78
/
check_cuda.py
executable file
·44 lines (36 loc) · 1.23 KB
/
check_cuda.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
#!/usr/bin/env python3
###################################################################################################
# Copyright (C) 2019-2023 Maxim Integrated Products, Inc. All Rights Reserved.
#
# Maxim Integrated Products, Inc. Default Copyright Notice:
# https://www.maximintegrated.com/en/aboutus/legal/copyrights.html
###################################################################################################
"""
Check whether PyTorch supports CUDA hardware acceleration.
"""
import signal
import sys
import torch
def signal_handler(
_signal,
_frame,
):
"""
Ctrl+C handler
"""
sys.exit(0)
if __name__ == '__main__':
signal.signal(signal.SIGINT, signal_handler)
print("System: ", sys.platform)
print("Python version: ", sys.version.replace('\n', ''))
print("PyTorch version: ", torch.__version__)
print("CUDA/ROCm acceleration: ", end='')
if not torch.cuda.is_available():
print("NOT available in PyTorch")
else:
print("available in PyTorch")
print("MPS acceleration: ", end='')
if not torch.backends.mps.is_available():
print("NOT available in PyTorch")
else:
print("available in PyTorch")