system/torch_check.py

import torch

if torch.cuda.is_available():
    print("CUDA is available.")
    print(f"GPU Name: {torch.cuda.get_device_name(0)}")
    print(f"CUDA Version: {torch.version.cuda}")
    print(f"cuDNN Version: {torch.backends.cudnn.version()}")
else:
    print("CUDA is not available. Check your driver and CUDA installation.")