File size: 438 Bytes
3c9c761 |
1 2 3 4 5 6 7 8 9 |
import torch, platform
print("torch:", torch.__version__)
print("torch compiled CUDA:", torch.version.cuda)
print("CUDA available:", torch.cuda.is_available())
if torch.cuda.is_available():
print("GPU:", torch.cuda.get_device_name(0))
print("Device capability (SM):", torch.cuda.get_device_capability(0)) # e.g., (9, 0)
print("Torch arch list:", torch.cuda.get_arch_list()) # what the wheel contains
|