import torch

if torch.cuda.is_available():
    print(f"CUDA is available. GPU: {torch.cuda.get_device_name(0)}")
    print(f"PyTorch CUDA version: {torch.version.cuda}")
else:
    print("CUDA is not available. PyTorch is running on CPU.")

# Check torchvision
try:
    import torchvision

    print(f"Torchvision version: {torchvision.__version__}")
    # You might need specific torchvision ops to fully check GPU integration,
    # but successful import and version check is a good start.
    print("Torchvision imported successfully.")
except ImportError:
    print("Torchvision could not be imported.")