torch.cuda用法

2022-10-04 13:36:13

下面是torch.cuda的一些常用函数:

import torch#检查cuda是否可以使用
torch.cuda.is_available()#查看当前gpu索引号
torch.cuda.current_device()#查看当前cuda流
torch.cuda.current_stream(device=0)#选择device
torch.cuda.device(1)#查看有多少个GPU设备
torch.cuda.device_count()#查看gpu的容量
torch.cuda.get_device_capability(device=0)
  • 作者:江南汪
  • 原文链接:https://blog.csdn.net/weixin_47156261/article/details/116717563
    更新时间:2022-10-04 13:36:13