diff --git a/recipes_source/torch_logs.py b/recipes_source/torch_logs.py index b5c3f0bd8a..45ca17584f 100644 --- a/recipes_source/torch_logs.py +++ b/recipes_source/torch_logs.py @@ -3,9 +3,7 @@ ========================================================================================== **Author:** `Michael Lazos `_ """ - import logging - ###################################################################### # # This tutorial introduces the ``TORCH_LOGS`` environment variable, as well as the Python API, and @@ -16,8 +14,6 @@ # This tutorial requires PyTorch 2.2.0 or later. # # - - ###################################################################### # Setup # ~~~~~~~~~~~~~~~~~~~~~ @@ -29,10 +25,22 @@ # There is also an environment variable ``TORCH_LOGS``, which can be used to # change logging settings at the command line. The equivalent environment # variable setting is shown for each example. - import torch -# exit cleanly if we are on a device that doesn't support torch.compile +###################################################################### +# Device Compatibility Check +# ~~~~~~~~~~~~~~~~~~~~~~~~~~~ +# The following code checks if the device supports torch.compile, which requires +# CUDA compute capability 7.0 or higher. +# +# .. important:: +# +# If your device does not meet this requirement, this tutorial will only print +# a skip message and exit. You will NOT see any of the logging output examples +# (Dynamo Tracing, Traced Graph, Fusion Decisions, Output Code) demonstrated +# below. To see the actual logging output, you need a CUDA-capable GPU with +# compute capability 7.0 or higher (e.g., V100, T4, or newer). +# if torch.cuda.get_device_capability() < (7, 0): print("Skipping because torch.compile is not supported on this device.") else: @@ -40,44 +48,33 @@ def fn(x, y): z = x + y return z + 2 - - 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() - - 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) - 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) - separator("") - ###################################################################### # Conclusion # ~~~~~~~~~~