Skip to content

Commit cd10048

Browse files
committed
docs: add fallback for torch.compile in torch_logs tutorial (#137285)
1 parent d54656f commit cd10048

File tree

1 file changed

+11
-6
lines changed

1 file changed

+11
-6
lines changed

recipes_source/torch_logs.py

Lines changed: 11 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -32,17 +32,22 @@
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("Skipping because torch.compile is not supported on this device.")
38-
else:
39-
@torch.compile()
35+
# Use torch.compile if supported, fallback to eager mode otherwise
36+
try:
37+
@torch.compile
4038
def fn(x, y):
4139
z = x + y
4240
return z + 2
41+
42+
inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
43+
except (RuntimeError, AssertionError):
44+
print("⚠️ torch.compile is not supported on this system. Falling back to eager mode.")
4345

46+
def fn(x, y):
47+
z = x + y
48+
return z + 2
4449

45-
inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
50+
inputs = (torch.ones(2, 2), torch.zeros(2, 2))
4651

4752

4853
# print separator and reset dynamo

0 commit comments

Comments
 (0)