Skip to content

Commit 1010779

Browse files
committed
Add cuda version check
1 parent 53db40d commit 1010779

File tree

1 file changed

+7
-0
lines changed

1 file changed

+7
-0
lines changed

recipes_source/torch_logs.py

Lines changed: 7 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,13 @@
3232

3333
import torch
3434

35+
# exit cleanly if we are on a device that doesn't support torch.compile
36+
if torch.cuda.get_device_capability() < (7, 0):
37+
print("Exiting because torch.compile is not supported on this device.")
38+
import sys
39+
40+
sys.exit(0)
41+
3542

3643
@torch.compile()
3744
def fn(x, y):

0 commit comments

Comments
 (0)