pytorch 验证GPU是否可用

最新推荐文章于 2025-06-10 19:30:00 发布

小猫不会去楼兰捉虫

最新推荐文章于 2025-06-10 19:30:00 发布

阅读量2.7k

点赞数 2

CC 4.0 BY-SA版权

分类专栏：深度学习 Python 文章标签： pytorch 深度学习机器学习

本文链接：https://blog.csdn.net/anglemanyi/article/details/124766488

深度学习同时被 2 个专栏收录

19 篇文章

订阅专栏

Python

17 篇文章

订阅专栏

该博客展示了如何在Python中使用PyTorch库来检测CUDA是否可用，并在GPU上运行简单的张量操作。通过`torch.cuda.is_available()`函数判断GPU状态，然后创建并打印在GPU上的张量，验证了GPU的运行环境。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

代码

import torch
flag = torch.cuda.is_available()
print(flag)

ngpu= 1
# Decide which device we want to run on
device = torch.device("cuda:0" if (torch.cuda.is_available() and ngpu > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())

运行结果如下表示GPU可用了

True
cuda:0
NVIDIA GeForce RTX 3060
tensor([[0.7567, 0.0458, 0.5618],
[0.6300, 0.2157, 0.8634],
[0.2069, 0.5930, 0.2276]], device='cuda:0')