import torch, platform
print("torch:", torch.__version__)
print("torch compiled CUDA:", torch.version.cuda)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
    print("GPU:", torch.cuda.get_device_name(0))
    print("Device capability (SM):", torch.cuda.get_device_capability(0))  # e.g., (9, 0)
    print("Torch arch list:", torch.cuda.get_arch_list())                  # what the wheel contains