首页 > 技术文章 > 验证cuda 及cudnn是否成功安装

Henry-ZHAO 2020-11-29 17:59 原文

import torch
from torch.backends import cudnn

if __name__== "__main__":
     print(torch.cuda.is_available())
     x =  torch.tensor([10.])
     x = x.cuda()
     print(cudnn.is_acceptable(x)) 
torch.backends.cudnn.enabled = True#启动cudnn
torch.backends.cudnn.benchmark = True#启动cudnn的不确定性算法的自动搜索优化模式,适合输入数据维度比较固定的情况,否则应该关闭此项。

推荐阅读