pythonmemory-managementgpunvidiapytorch

How do I check if PyTorch is using the GPU?


How do I check if PyTorch is using the GPU? The nvidia-smi command can detect GPU activity, but I want to check it directly from inside a Python script.


Solution

  • These functions should help:

    >>> import torch
    
    >>> torch.cuda.is_available()
    True
    
    >>> torch.cuda.device_count()
    1
    
    >>> torch.cuda.current_device()
    0
    
    >>> torch.cuda.device(0)
    <torch.cuda.device at 0x7efce0b03be0>
    
    >>> torch.cuda.get_device_name(0)
    'GeForce GTX 950M'
    

    This tells us: