import torch

# Check if GPU is available
if torch.cuda.is_available():
    device = torch.device("cuda")
    
    # Optional: Print GPU model name for confirmation
    gpu_name = torch.cuda.get_device_name(0)
    print("="*50)
    print(f"✅ GPU Detected and Available! Using device: {device}")
    print(f"   GPU Model: {gpu_name}")
    print("="*50)
else:
    device = torch.device("cpu")
    print("="*50)
    print("⚠️ WARNING: CUDA not available. Falling back to CPU.")
    print("   (If you have a GPU, ensure you installed the correct PyTorch version for CUDA.)")
    print("="*50)