如何检查 PyTorch 是否正在使用 GPU？

如何检查 PyTorch 是否正在使用 GPU？这nvidia-smi命令可以检测 GPU 活动，但我想直接从 Python 脚本内部检查它。

这些功能应该有助于：

>>> 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'

这告诉我们：

• CUDA 可用并且可由一台设备使用。
• Device 0指的是GPUGeForce GTX 950M，目前由 PyTorch 选择。