Modify example code to fix torch logs doc

Author: nluu175

recipes_source/torch_logs.py

@@ -31,52 +31,65 @@
 #   variable setting is shown for each example.
 
 import torch
+import sys
 
-# 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
-
-
-    inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
 
+def env_setup():
+    """Set up environment for running the example. Exit cleanly if CUDA is not available."""
+    if not torch.cuda.is_available():
+        print("CUDA is not available. Exiting.")
+        sys.exit(0)
+        
+    if torch.cuda.get_device_capability() < (7, 0):
+        print("Skipping because torch.compile is not supported on this device.")
+        sys.exit(0)
 
-# print separator and reset dynamo
-# between each example
-    def separator(name):
-        print(f"==================={name}=========================")
-        torch._dynamo.reset()
 
+def separator(name):
+    """Print separator and reset dynamo between each example"""
+    print(f"\n{'='*20} {name} {'='*20}")
+    torch._dynamo.reset()
 
-    separator("Dynamo Tracing")
-# View dynamo tracing
-# TORCH_LOGS="+dynamo"
-    torch._logging.set_logs(dynamo=logging.DEBUG)
-    fn(*inputs)
 
-    separator("Traced Graph")
-# View traced graph
-# TORCH_LOGS="graph"
-    torch._logging.set_logs(graph=True)
-    fn(*inputs)
+def run_debugging_suite():
+    """Run the complete debugging suite with all logging options"""
+    env_setup()
 
-    separator("Fusion Decisions")
-# View fusion decisions
-# TORCH_LOGS="fusion"
-    torch._logging.set_logs(fusion=True)
-    fn(*inputs)
-
-    separator("Output Code")
-# View output code generated by inductor
-# TORCH_LOGS="output_code"
-    torch._logging.set_logs(output_code=True)
-    fn(*inputs)
+    @torch.compile()
+    def fn(x, y):
+        z = x + y
+        return z + 2
 
-    separator("")
+    inputs = (
+        torch.ones(2, 2, device="cuda"),
+        torch.zeros(2, 2, device="cuda")
+    )
+    
+    logging_scenarios = [
+        # View dynamo tracing; TORCH_LOGS="+dynamo"
+        ("Dynamo Tracing", {"dynamo": logging.DEBUG}), 
+
+        # View traced graph; TORCH_LOGS="graph"
+        ("Traced Graph", {"graph": True}), 
+        
+        # View fusion decisions; TORCH_LOGS="fusion"
+        ("Fusion Decisions", {"fusion": True}), 
+        
+        # View output code generated by inductor; TORCH_LOGS="output_code"
+        ("Output Code", {"output_code": True}) 
+    ]
+    
+    for name, log_config in logging_scenarios:
+        separator(name)
+        torch._logging.set_logs(**log_config)
+        try:
+            result = fn(*inputs)
+            print(f"Function output shape: {result.shape}")
+        except Exception as e:
+            print(f"Error during {name}: {str(e)}")
+
+if __name__ == "__main__":
+    run_debugging_suite()
 
 ######################################################################
 # Conclusion