/ src / liger_kernel / env_report.py
env_report.py
 1  import platform
 2  import sys
 3  
 4  
 5  def print_env_report():
 6      """
 7  
 8      Prints a report of the environment. Useful for debugging and reproducibility.
 9      Usage:
10      ```
11      python -m liger_kernel.env_report
12      ```
13  
14      """
15      print("Environment Report:")
16      print("-------------------")
17      print(f"Operating System: {platform.platform()}")
18      print(f"Python version: {sys.version.split()[0]}")
19  
20      try:
21          import torch
22  
23          print(f"PyTorch version: {torch.__version__}")
24          cuda_version = (
25              torch.version.cuda if torch.cuda.is_available() else "Not available"
26          )
27          print(f"CUDA version: {cuda_version}")
28      except ImportError:
29          print("PyTorch: Not installed")
30          print("CUDA version: Unable to query")
31  
32      try:
33          import triton
34  
35          print(f"Triton version: {triton.__version__}")
36      except ImportError:
37          print("Triton: Not installed")
38  
39      try:
40          import transformers
41  
42          print(f"Transformers version: {transformers.__version__}")
43      except ImportError:
44          print("Transformers: Not installed")
45  
46  
47  if __name__ == "__main__":
48      print_env_report()