Fix torch_logs tutorial device gate and CPU fallback (#137285)

Author: yyouretoast

recipes_source/torch_logs.py

@@ -32,51 +32,63 @@
 
 import torch
 
-# exit cleanly if we are on a device that doesn't support torch.compile
-if torch.cuda.get_device_capability() < (7, 0):
-    print("Skipping because torch.compile is not supported on this device.")
-else:
-    @torch.compile()
-    def fn(x, y):
-        z = x + y
-        return z + 2
+# Determine device: CUDA if supported, otherwise CPU
+device = "cuda" if torch.cuda.is_available() and torch.cuda.get_device_capability() >= (7, 0) else "cpu"
 
+# Verify if torch.compile is supported in the current environment
+try:
+    torch.compile(lambda x: x + 1)(torch.ones(1, device=device))
+    compile_supported = True
+except Exception:
+    compile_supported = False
 
-    inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
+if not compile_supported:
+    print("Skipping because torch.compile is not supported on this device/environment.")
+    import sys
+    sys.exit(0)
+
+
+@torch.compile()
+def fn(x, y):
+    z = x + y
+    return z + 2
+
+
+inputs = (torch.ones(2, 2, device=device), torch.zeros(2, 2, device=device))
 
 
 # print separator and reset dynamo
 # between each example
-    def separator(name):
-        print(f"==================={name}=========================")
-        torch._dynamo.reset()
+def separator(name):
+    print(f"==================={name}=========================")
+    torch._dynamo.reset()
 
 
-    separator("Dynamo Tracing")
+separator("Dynamo Tracing")
 # View dynamo tracing
 # TORCH_LOGS="+dynamo"
-    torch._logging.set_logs(dynamo=logging.DEBUG)
-    fn(*inputs)
+torch._logging.set_logs(dynamo=logging.DEBUG)
+fn(*inputs)
 
-    separator("Traced Graph")
+separator("Traced Graph")
 # View traced graph
 # TORCH_LOGS="graph"
-    torch._logging.set_logs(graph=True)
-    fn(*inputs)
+torch._logging.set_logs(graph=True)
+fn(*inputs)
 
-    separator("Fusion Decisions")
+separator("Fusion Decisions")
 # View fusion decisions
 # TORCH_LOGS="fusion"
-    torch._logging.set_logs(fusion=True)
-    fn(*inputs)
+torch._logging.set_logs(fusion=True)
+fn(*inputs)
 
-    separator("Output Code")
+separator("Output Code")
 # View output code generated by inductor
 # TORCH_LOGS="output_code"
-    torch._logging.set_logs(output_code=True)
-    fn(*inputs)
+torch._logging.set_logs(output_code=True)
+fn(*inputs)
 
-    separator("")
+separator("")
 
 ######################################################################
 # Conclusion