pytorch

How to Check If PyTorch Is Using NVIDIA GPU for CUDA Acceleration

PyTorch is an Artificial Intelligence/Machine Learning/Deep Learning framework from Meta (the company behind Facebook). PyTorch can run on CPU as well as NVIDIA and AMD GPUs. Training PyTorch models on GPUs is faster compared to training them on CPU.

In this article, we will show you how to check if PyTorch is using NVIDIA GPU/CUDA to accelerate the training of your PyTorch AI/ML/DL models faster.

Topic of Contents:

  1. Opening the Python Interactive Interpreter/Shell
  2. Importing PyTorch on the Python Interactive Interpreter/Shell
  3. Checking If PyTorch with CUDA Support Is Installed
  4. Checking If PyTorch Can Use NVIDIA GPU for CUDA Acceleration
  5. Checking the Number of CUDA Capable Devices Available for PyTorch
  6. Checking the Currently Used NVIDIA GPU of PyTorch
  7. Conclusion

Opening the Python Interactive Interpreter/Shell

You can check if PyTorch is using NVIDIA GPU/CUDA to accelerate the PyTorch AI/ML/DL model training from the Python interactive interpreter/shell.

To open the Python interactive interpreter/shell, run one of the following commands (depending on how Python is configured on your system) from a Terminal program:

$ python

or

$ python3

The Python interactive interpreter/shell should be opened.

Importing PyTorch on the Python Interactive Interpreter/Shell

To import PyTorch on the Python interactive interpreter/shell, run the following line of code:

$ import torch

Checking If PyTorch with CUDA Support Is Installed

Once PyTorch is imported, check the version of PyTorch that you installed on your computer to verify if it was built to use the NVIDIA GPU/CUDA acceleration.

To check the version of PyTorch that you installed, run the following line of code:

torch.__version__

As you can see, we have PyTorch 2.0.1 installed on our computer and it was built to use the NVIDIA CUDA 11.8 (cu118) to accelerate the training of PyTorch models.

A black screen with white text Description automatically generated

Checking If PyTorch Can Use NVIDIA GPU for CUDA Acceleration

To check whether the NVIDIA CUDA acceleration is available for PyTorch, run the following line of code:

$ torch.cuda.is_available()

If the NVIDIA CUDA acceleration is available for PyTorch, “True” will be printed on the screen. Otherwise, “False” will be printed.

Checking the Number of CUDA Capable Devices Available for PyTorch

To find out the number of NVIDIA GPUs capable of CUDA acceleration (to speed up the PyTorch training), run the following line of code:

$ torch.cuda.device_count()

As you can see, we have two CUDA capable NVIDIA GPUs installed on our computer (that PyTorch can use to speed up the training of AI/ML/DL models).

Python indices start from zero. So, to find the name/model of the first CUDA-capable NVIDIA GPU of your computer, run the following line of code:

$ torch.cuda.get_device_name(0)

As you can see, the first CUDA-capable NVIDIA GPU installed on our computer is NVIDIA GeForce RTX 4070.

In the same way, you can find the name/model of the second CUDA-capable NVIDIA GPU of your computer with the following line of code:

$ torch.cuda.get_device_name(1)

As you can see, the second CUDA-capable NVIDIA GPU installed on our computer is NVIDIA GeForce GTX 1050 Ti.

Checking the Currently Used NVIDIA GPU of PyTorch

To find the NVIDIA GPU that PyTorch is currently using for CUDA acceleration, run the following line of code:

$ torch.cuda.current_device()

The index/ID of the NVIDIA GPU should be printed.

To resolve the index/ID of PyTorch’s currently used GPU to the name/model of the GPU, run the following line of code:

$ torch.cuda.get_device_name(0)

As you can see, PyTorch is currently using an NVIDIA GeForce RTX 4070 GPU for CUDA acceleration.

If you’re only interested in the name/model of the GPU that PyTorch is currently using instead of the index/ID number of the currently used GPU, you can run the following line of code. You will see the same result.

$ torch.cuda.get_device_name(torch.cuda.current_device())

Conclusion

In this article, we showed you how to check if the installed version of PyTorch can use the NVIDIA GPUs for CUDA acceleration. We also showed you how to check if you have CUDA-capable NVIDIA GPUs available on your computer and if PyTorch can use them to accelerate the training of PyTorch AI/ML/DL models. We also showed you how to check the number of CUDA-capable NVIDIA GPUs available in PyTorch and list them by name/model as well. Finally, we showed you how to check the currently used NVIDIA GPU of PyTorch for CUDA acceleration.

About the author

Shahriar Shovon

Freelancer & Linux System Administrator. Also loves Web API development with Node.js and JavaScript. I was born in Bangladesh. I am currently studying Electronics and Communication Engineering at Khulna University of Engineering & Technology (KUET), one of the demanding public engineering universities of Bangladesh.