Nexa_Labs / scripts /check_cuda.py
Allanatrix's picture
Upload 57 files
d8328bf verified
"""Quick script to check CUDA availability and PyTorch installation."""
import sys
print("=" * 80)
print("CUDA/PyTorch GPU Check")
print("=" * 80)
try:
import torch
print(f"βœ“ PyTorch version: {torch.__version__}")
cuda_available = torch.cuda.is_available()
print(f"βœ“ CUDA available: {cuda_available}")
if cuda_available:
print(f"βœ“ CUDA version: {torch.version.cuda}")
print(f"βœ“ cuDNN version: {torch.backends.cudnn.version()}")
print(f"βœ“ GPU count: {torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
print(f" GPU {i}: {torch.cuda.get_device_name(i)}")
props = torch.cuda.get_device_properties(i)
print(f" Memory: {props.total_memory / (1024**3):.1f} GB")
else:
print("\n⚠️ CUDA is NOT available!")
print("\nPossible causes:")
print("1. PyTorch was installed without CUDA support")
print("2. NVIDIA drivers not installed")
print("3. CUDA toolkit not installed")
print("\nTo fix:")
print("1. Check NVIDIA drivers: nvidia-smi")
print("2. Reinstall PyTorch with CUDA:")
print(" pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121")
print(" (adjust cu121 to match your CUDA version)")
except ImportError:
print("βœ— PyTorch not installed")
sys.exit(1)
print("=" * 80)