Advertisements
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.