[python] How to check if pytorch is using the GPU?

As it hasn't been proposed here, I'm adding a method using torch.device, as this is quite handy, also when initializing tensors on the correct device.

# setting device on GPU if available, else CPU
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
print('Using device:', device)
print()

#Additional Info when using cuda
if device.type == 'cuda':
    print(torch.cuda.get_device_name(0))
    print('Memory Usage:')
    print('Allocated:', round(torch.cuda.memory_allocated(0)/1024**3,1), 'GB')
    print('Cached:   ', round(torch.cuda.memory_reserved(0)/1024**3,1), 'GB')

Edit: torch.cuda.memory_cached has been renamed to torch.cuda.memory_reserved. So use memory_cached for older versions.

Output:

Using device: cuda

Tesla K80
Memory Usage:
Allocated: 0.3 GB
Cached:    0.6 GB

As mentioned above, using device it is possible to:

  • To move tensors to the respective device:

      torch.rand(10).to(device)
    
  • To create a tensor directly on the device:

      torch.rand(10, device=device)
    

Which makes switching between CPU and GPU comfortable without changing the actual code.


Edit:

As there has been some questions and confusion about the cached and allocated memory I'm adding some additional information about it:


You can either directly hand over a device as specified further above in the post or you can leave it None and it will use the current_device().


Additional note: Old graphic cards with Cuda compute capability 3.0 or lower may be visible but cannot be used by Pytorch!
Thanks to hekimgil for pointing this out! - "Found GPU0 GeForce GT 750M which is of cuda capability 3.0. PyTorch no longer supports this GPU because it is too old. The minimum cuda capability that we support is 3.5."

Examples related to python

programming a servo thru a barometer Is there a way to view two blocks of code from the same file simultaneously in Sublime Text? python variable NameError Why my regexp for hyphenated words doesn't work? Comparing a variable with a string python not working when redirecting from bash script is it possible to add colors to python output? Get Public URL for File - Google Cloud Storage - App Engine (Python) Real time face detection OpenCV, Python xlrd.biffh.XLRDError: Excel xlsx file; not supported Could not load dynamic library 'cudart64_101.dll' on tensorflow CPU-only installation

Examples related to memory-management

When to create variables (memory management) How to check if pytorch is using the GPU? How to delete multiple pandas (python) dataframes from memory to save RAM? Is there a way to delete created variables, functions, etc from the memory of the interpreter? C++ error : terminate called after throwing an instance of 'std::bad_alloc' How to delete object? Android Studio - How to increase Allocated Heap Size Implementing IDisposable correctly Calculating Page Table Size Pointer-to-pointer dynamic two-dimensional array

Examples related to gpu

How do I use TensorFlow GPU? How to check if pytorch is using the GPU? NVIDIA NVML Driver/library version mismatch NVIDIA-SMI has failed because it couldn't communicate with the NVIDIA driver How to get current available GPUs in tensorflow? How to tell if tensorflow is using gpu acceleration from inside python shell? Using Java with Nvidia GPUs (CUDA) Is it possible to run CUDA on AMD GPUs? How do I choose grid and block dimensions for CUDA kernels? Can I run CUDA on Intel's integrated graphics processor?

Examples related to nvidia

How to check if pytorch is using the GPU? NVIDIA NVML Driver/library version mismatch How do I select which GPU to run a job on? Error Message : Cannot find or open the PDB file Is it possible to run CUDA on AMD GPUs? Difference between nVidia Quadro and Geforce cards? How do I choose grid and block dimensions for CUDA kernels? Understanding CUDA grid dimensions, block dimensions and threads organization (simple explanation)

Examples related to pytorch

Pytorch tensor to numpy array How to initialize weights in PyTorch? How to check if pytorch is using the GPU? PyTorch: How to get the shape of a Tensor as a list of int Pytorch reshape tensor dimension Best way to save a trained model in PyTorch? Model summary in pytorch How does the "view" method work in PyTorch?