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

To use a GPU with Python and PyTorch, you will first need to install the appropriate GPU drivers and CUDA Toolkit for your GPU. Then, you can install PyTorch with the GPU support option.

Once PyTorch is installed, you can use the .cuda() method to transfer data and models to the GPU for accelerated processing. For example:

import torch

# Create a tensor on the CPU
x = torch.randn(3, 3)

# Transfer it to the GPU
x = x.cuda()

When you create a model, you can also specify that it should be created on the GPU by setting the device argument to cuda. For example:

import torch.nn as nn

model = nn.Linear(3, 3).to(device='cuda')

You can also use the torch.cuda.is_available() function to check if a GPU is available.

For more information about using GPUs with PyTorch, please see the PyTorch documentation.

Edit this code on GitHub