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.