Skip to content

Commit 6cc87a6

Browse files
authored
Merge pull request #1 from zzadxz/zzadxz-patch-1
Fix torch_logs to run on CPU devices
2 parents 67ec2a5 + 6f4b550 commit 6cc87a6

File tree

1 file changed

+21
-25
lines changed

1 file changed

+21
-25
lines changed

recipes_source/torch_logs.py

Lines changed: 21 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -32,51 +32,47 @@
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+
@torch.compile()
36+
def fn(x, y):
37+
z = x + y
38+
return z + 2
4339

4440

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

4743

4844
# print separator and reset dynamo
4945
# between each example
50-
def separator(name):
51-
print(f"==================={name}=========================")
52-
torch._dynamo.reset()
46+
def separator(name):
47+
print(f"==================={name}=========================")
48+
torch._dynamo.reset()
5349

5450

55-
separator("Dynamo Tracing")
51+
separator("Dynamo Tracing")
5652
# View dynamo tracing
5753
# TORCH_LOGS="+dynamo"
58-
torch._logging.set_logs(dynamo=logging.DEBUG)
59-
fn(*inputs)
54+
torch._logging.set_logs(dynamo=logging.DEBUG)
55+
fn(*inputs)
6056

61-
separator("Traced Graph")
57+
separator("Traced Graph")
6258
# View traced graph
6359
# TORCH_LOGS="graph"
64-
torch._logging.set_logs(graph=True)
65-
fn(*inputs)
60+
torch._logging.set_logs(graph=True)
61+
fn(*inputs)
6662

67-
separator("Fusion Decisions")
63+
separator("Fusion Decisions")
6864
# View fusion decisions
6965
# TORCH_LOGS="fusion"
70-
torch._logging.set_logs(fusion=True)
71-
fn(*inputs)
66+
torch._logging.set_logs(fusion=True)
67+
fn(*inputs)
7268

73-
separator("Output Code")
69+
separator("Output Code")
7470
# View output code generated by inductor
7571
# TORCH_LOGS="output_code"
76-
torch._logging.set_logs(output_code=True)
77-
fn(*inputs)
72+
torch._logging.set_logs(output_code=True)
73+
fn(*inputs)
7874

79-
separator("")
75+
separator("")
8076

8177
######################################################################
8278
# Conclusion

0 commit comments

Comments
 (0)