From 1261186cb7b298f66d6155c2abac337664743f18 Mon Sep 17 00:00:00 2001
From: Michael Lazos <mlazos@fb.com>
Date: Tue, 30 Jan 2024 22:11:10 +0000
Subject: [PATCH] Exit cleanly

---
 recipes_source/torch_logs.py | 50 +++++++++++++++++-------------------
 1 file changed, 23 insertions(+), 27 deletions(-)

diff --git a/recipes_source/torch_logs.py b/recipes_source/torch_logs.py
index 86bacda5503..90edf4c5031 100644
--- a/recipes_source/torch_logs.py
+++ b/recipes_source/torch_logs.py
@@ -34,53 +34,49 @@
 
 # exit cleanly if we are on a device that doesn't support torch.compile
 if torch.cuda.get_device_capability() < (7, 0):
-    print("Exiting because torch.compile is not supported on this device.")
-    import sys
+    print("Skipping because torch.compile is not supported on this device.")
+else:
+    @torch.compile()
+    def fn(x, y):
+        z = x + y
+        return z + 2
 
-    sys.exit(0)
 
-
-@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"))
+    inputs = (torch.ones(2, 2, device="cuda"), torch.zeros(2, 2, device="cuda"))
 
 
 # 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