当前位置:网站首页>Tensorflow and torch code verify whether CUDA is successfully installed

Tensorflow and torch code verify whether CUDA is successfully installed

2022-07-06 19:05:00 The mountain of ignorance, the valley of despair, the slope of

tensorflow Code validation cuda Is the installation successful
tensorflow test cuda Is the installation successful , Test code environment

tensorflow

import tensorflow as tf

tf.test.is_gpu_available()

perhaps

sess = tf.Session(config=tf.ConfigProto(log_device_placement=True))

torch

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

gpunumber = 1

Decide which device we want to run on

device = torch.device("cuda:0" if (torch.cuda.is_available() and gpunumber > 0) else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())
原网站

版权声明
本文为[The mountain of ignorance, the valley of despair, the slope of ]所创,转载请带上原文链接,感谢
https://yzsam.com/2022/187/202207061110146671.html

随机推荐