torch-is-cuda-available.py

   1 import torch
   2
   3 # Check if GPU is available
   4 if torch.cuda.is_available():
   5     device = torch.device("cuda")
   6
   7     # Optional: Print GPU model name for confirmation
   8     gpu_name = torch.cuda.get_device_name(0)
   9     print("="*50)
  10     print(f"✅ GPU Detected and Available! Using device: {device}")
  11     print(f"   GPU Model: {gpu_name}")
  12     print("="*50)
  13 else:
  14     device = torch.device("cpu")
  15     print("="*50)
  16     print("⚠️ WARNING: CUDA not available. Falling back to CPU.")
  17     print("   (If you have a GPU, ensure you installed the correct PyTorch version for CUDA.)")
  18     print("="*50)