File size: 881 Bytes
75ec748 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 |
import torch
def test_cuda():
print("Checking CUDA availability...")
cuda_available = torch.cuda.is_available()
print(f"CUDA Available: {cuda_available}")
if cuda_available:
print(f"CUDA Version: {torch.version.cuda}")
print(f"PyTorch version: {torch.__version__}")
print(f"Device count: {torch.cuda.device_count()}")
for i in range(torch.cuda.device_count()):
print(f"Device {i}: {torch.cuda.get_device_name(i)}")
# Simple tensor test
print("Running a simple tensor operation on CUDA...")
x = torch.tensor([1.0, 2.0, 3.0], device="cuda")
y = torch.tensor([4.0, 5.0, 6.0], device="cuda")
z = x + y
print(f"Result (on CUDA): {z}")
else:
print("CUDA is not available on this system.")
if __name__ == "__main__":
test_cuda()
|