How to verify your CUDA and pytorch

Do things2024-10-26 16:34

1.Test CUDA

This script will attempt to create a tensor on the GPU and perform a simple computation.

python 复制代码

import torch

def test_cuda():
    # Check CUDA availability
    if torch.cuda.is_available():
        print("CUDA is available on this device.")
        # Set device to CUDA
        device = torch.device("cuda")
        # Create a tensor and move it to CUDA
        x = torch.rand(10, 10).to(device)
        print("Successfully created a tensor on CUDA:", x)
        # Perform a simple computation
        y = x * x
        print("Output of the computation:", y)
    else:
        print("CUDA is not available. Check your installation and driver.")

if __name__ == "__main__":
    test_cuda()

Run the script: Use a terminal or command prompt to run the script with Python.

bash 复制代码

python test_cuda.py

The script should print whether CUDA is available, show the tensor created on the GPU, and display the output of a simple multiplication operation. If there are any errors during these steps, they will help pinpoint what might be wrong with your CUDA setup.

2.Check PyTorch Compatibility

python 复制代码

# test_torch.py
import torch
print(torch.__version__)
print(torch.cuda.is_available())

Run the script: Use a terminal or command prompt to run the script with Python.

bash 复制代码

python test_torch.py

If torch.cuda.is_available() returns True, then PyTorch is able to use CUDA.