pytorch


How to check if PyTorch is using the GPU?

The following basic code, will import PyTorch into a project and test if GPU capabilities are enabled and available.

import torch

# Should produce "True"
torch.cuda.is_available()

# Should produce the number of available devices, if you have one device it should produce the value "1"
torch.cuda.device_count()

# If there is a device and it is the first one, it should produce the "0"
torch.cuda.current_device()

# Assuming that there is at least one device, with the following two commands we will get some information on the first available device

# Should produce something similar to "<torch.cuda.device object at 0x7f12b1a298d0>"
torch.cuda.device(0)

# Should produce something similar to "'GeForce GTX 1050 Ti'"
torch.cuda.get_device_name(0)

We are using Ubuntu 20.04 LTS and the NVidia drivers were installed automatically during installation.