Skip to content

Commit 2f00fa4

Browse files
committed
Fix torch_logs tutorial to support CPU-only environments
Remove the hard CUDA gate that caused the tutorial to render as 'Skipped' on the PyTorch docs website. torch.compile works on both CUDA (via Triton) and CPU (via C++ Inductor), so no guard or early exit is needed. Changes: - Use CUDA if available and compute capability >= 7.0, else fall back to CPU (prevents crashes on older CUDA devices while supporting CPU execution). - Remove try/except compilation probe (unnecessary). - Remove sys.exit(0) early exit (dangerous in Sphinx-Gallery subprocess mode). - De-indent all tutorial code from the original if/else block. Fixes pytorch/pytorch#137285
1 parent 3cdec7b commit 2f00fa4

1 file changed

Lines changed: 26 additions & 25 deletions

File tree

recipes_source/torch_logs.py

Lines changed: 26 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -32,51 +32,52 @@
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()
40-
def fn(x, y):
41-
z = x + y
42-
return z + 2
35+
# Determine device: use CUDA if available and compute capability >= 7.0, otherwise fall back to CPU.
36+
# torch.compile works on both CUDA (via Triton) and CPU (via C++ Inductor).
37+
device = "cuda" if torch.cuda.is_available() and torch.cuda.get_device_capability() >= (7, 0) else "cpu"
4338

4439

45-
inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
40+
@torch.compile()
41+
def fn(x, y):
42+
z = x + y
43+
return z + 2
44+
45+
46+
inputs = (torch.ones(2, 2, device=device), torch.zeros(2, 2, device=device))
4647

4748

4849
# print separator and reset dynamo
4950
# between each example
50-
def separator(name):
51-
print(f"==================={name}=========================")
52-
torch._dynamo.reset()
51+
def separator(name):
52+
print(f"==================={name}=========================")
53+
torch._dynamo.reset()
5354

5455

55-
separator("Dynamo Tracing")
56+
separator("Dynamo Tracing")
5657
# View dynamo tracing
5758
# TORCH_LOGS="+dynamo"
58-
torch._logging.set_logs(dynamo=logging.DEBUG)
59-
fn(*inputs)
59+
torch._logging.set_logs(dynamo=logging.DEBUG)
60+
fn(*inputs)
6061

61-
separator("Traced Graph")
62+
separator("Traced Graph")
6263
# View traced graph
6364
# TORCH_LOGS="graph"
64-
torch._logging.set_logs(graph=True)
65-
fn(*inputs)
65+
torch._logging.set_logs(graph=True)
66+
fn(*inputs)
6667

67-
separator("Fusion Decisions")
68+
separator("Fusion Decisions")
6869
# View fusion decisions
6970
# TORCH_LOGS="fusion"
70-
torch._logging.set_logs(fusion=True)
71-
fn(*inputs)
71+
torch._logging.set_logs(fusion=True)
72+
fn(*inputs)
7273

73-
separator("Output Code")
74+
separator("Output Code")
7475
# View output code generated by inductor
7576
# TORCH_LOGS="output_code"
76-
torch._logging.set_logs(output_code=True)
77-
fn(*inputs)
77+
torch._logging.set_logs(output_code=True)
78+
fn(*inputs)
7879

79-
separator("")
80+
separator("")
8081

8182
######################################################################
8283
# Conclusion

0 commit comments

Comments
 (0)