19 lines
608 B
Python
19 lines
608 B
Python
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.")
|