python-pytorchHow can I use a GPU to run Python and PyTorch?

Using a GPU to run Python and PyTorch is relatively straightforward. Here is an example of code that can be used to check if a GPU is being used:

import torch

if torch.cuda.is_available():
    device = torch.device("cuda")  # a CUDA device object
    print("CUDA is available")
else:
    print("CUDA is not available")

Output example

CUDA is available

To use a GPU with PyTorch, the code needs to be modified to include the following components:

1. A device object that will be used to run operations on the GPU. This can be done with the following line of code: device = torch.device("cuda").

2. A function that will transfer data from the CPU to the GPU. This can be done with the following line of code: data = data.to(device).

3. A function that will transfer data from the GPU to the CPU. This can be done with the following line of code: data = data.cpu().

4. A function that will run operations on the GPU. This can be done with the following line of code: output = model(data).

Helpful links

• https://pytorch.org/tutorials/beginner/blitz/cifar10_tutorial.html
• https://pytorch.org/docs/stable/notes/cuda.html

Edit this code on GitHub